Merge remote-tracking branch 'origin/main' into litellm-providers/price-sync-aws-bedrock

.github/workflows/test-linting.yml

@@ -524,6 +524,12 @@ jobs:
       - name: check_provider_folders_documented
         run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py
 
+      - name: check_provider_endpoints_support_generated
+        run: uv run --no-sync python scripts/generate_provider_endpoints_support.py --check
+
+      - name: test_generate_provider_endpoints_support
+        run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_generate_provider_endpoints_support.py
+
       - name: check_prisma_binary_cache
         run: uv run --no-sync python ./tests/code_coverage_tests/check_prisma_binary_cache.py
 

litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py

@@ -0,0 +1,639 @@
+"""Read-only checks that the rows already in the database can take the pending schema
+changes, run before `prisma migrate deploy` and before each request-log index build.
+
+Postgres has a short list of schema statements that fail because of existing rows: a
+btree index whose entry for some row would not fit in a third of a page (2704 bytes on
+the default 8 KB page, counting every indexed column), UNIQUE (duplicate values), SET NOT NULL and ADD COLUMN NOT NULL with no default (nulls, any row at all),
+FOREIGN KEY (orphans) and CHECK (rows failing the expression). Each pending statement
+of one of those kinds becomes one count query derived from the statement text and the
+column types in the catalog. A count above zero names the table, the column, the
+statement it blocks and the SQL that clears it, and the migration is refused before
+anything is applied. Every other statement a migration can contain cannot fail on
+existing data and gets no query. A statement the derivation cannot read gets none
+either, so the migration keeps today's behavior for it.
+"""
+
+import os
+import re
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from functools import reduce
+from itertools import accumulate
+from pathlib import Path
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, Literal
+
+from litellm_proxy_extras._logging import logger
+
+if TYPE_CHECKING:
+    import psycopg
+    from psycopg import sql
+
+DEFAULT_BLOCK_SIZE: Final = 8192
+INDEX_TUPLE_HEADER_BYTES: Final = 8
+NULL_BITMAP_BYTES: Final = 8
+VARLENA_HEADER_BYTES: Final = 4
+SHORT_VARLENA_MAX_BYTES: Final = 126
+_PAGE_HEADER_WITH_THREE_LINE_POINTERS_BYTES: Final = 40
+_BTREE_PAGE_OPAQUE_BYTES: Final = 16
+_HEAP_TID_BYTES: Final = 8
+_ALIGNMENT: Final = MappingProxyType({"c": 1, "s": 2, "i": 4, "d": 8})
+_OCTET_LENGTH_TYPES: Final = frozenset({"text", "varchar", "bpchar", "bytea", "name"})
+MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIMEOUT"
+DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0
+REFUSAL: Final = "Existing rows would make a pending schema change fail:"
+LEGACY_SPEND_LOG_KEY_REMEDY: Final = (
+    "Rows written before v1.99.0 hold the raw key. Hash them in place (one-way, idempotent):\n"
+    '    UPDATE "LiteLLM_SpendLogs" SET api_key = encode(sha256(convert_to('
+    "regexp_replace(api_key, '^[Bb][Ee][Aa][Rr][Ee][Rr] ', ''), 'UTF8')), 'hex')\n"
+    "    WHERE api_key <> '' AND api_key !~ '^[0-9a-fA-F]{64}$' AND api_key !~ '^hashed-jwt-[0-9A-Fa-f]{64}$';"
+)
+
+Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"]
+
+
+@dataclass(frozen=True, slots=True)
+class Check:
+    kind: Kind
+    table: str
+    columns: tuple[str, ...]
+    statement: str
+    source: str
+    index: str | None = None
+    referenced_table: str | None = None
+    referenced_columns: tuple[str, ...] = ()
+    expression: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class BlockedChange:
+    check: Check
+    rows: int
+    detail: str
+
+
+@dataclass(frozen=True, slots=True)
+class CouldNotVerify:
+    check: Check
+    timeout_seconds: float
+
+
+Finding = BlockedChange | CouldNotVerify
+
+
+@dataclass(frozen=True, slots=True)
+class ColumnType:
+    """One column as `pg_attribute` and `pg_type` describe it: `length` is `attlen` (-1 for a
+    variable-length value), `align` is `attalign` (c, s, i or d), `typmod` is `atttypmod`."""
+
+    type_name: str
+    length: int
+    align: str
+    typmod: int
+
+    @property
+    def is_variable_length(self) -> bool:
+        return self.length == -1
+
+    @property
+    def can_overflow_btree(self) -> bool:
+        """Whether one value of this column alone can push a btree entry past the limit: a
+        variable-length type with no declared cap, or a cap of more bytes than the limit."""
+        if not self.is_variable_length:
+            return False
+        return self.typmod == -1 or (self.typmod - VARLENA_HEADER_BYTES) * 4 > btree_max_entry_bytes(DEFAULT_BLOCK_SIZE)
+
+
+ColumnTypeLookup = Callable[[str, str], ColumnType | None]
+
+_IDENT: Final = r'"?(\w+)"?'
+_QUALIFIED: Final = r'(?:"?\w+"?\s*\.\s*)?"?(\w+)"?'
+_CREATE_INDEX: Final = re.compile(
+    rf"^\s*CREATE\s+(UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?{_IDENT}\s+ON\s+(?:ONLY\s+)?"
+    rf"{_QUALIFIED}\s*(?:USING\s+(\w+)\s*)?\((.*?)\)",
+    re.IGNORECASE | re.DOTALL,
+)
+_ALTER_TABLE: Final = re.compile(
+    rf"^\s*ALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?{_QUALIFIED}\s+(.*)$", re.IGNORECASE | re.DOTALL
+)
+_SET_NOT_NULL: Final = re.compile(rf"^ALTER\s+(?:COLUMN\s+)?{_IDENT}\s+SET\s+NOT\s+NULL$", re.IGNORECASE)
+_ADD_COLUMN: Final = re.compile(
+    rf"^ADD\s+(?:COLUMN\s+)?(?:IF\s+NOT\s+EXISTS\s+)?(?!(?:CONSTRAINT|FOREIGN|UNIQUE|CHECK|PRIMARY|EXCLUDE)\b){_IDENT}\s+(.*)$",
+    re.IGNORECASE | re.DOTALL,
+)
+_CONSTRAINT: Final = rf"^ADD\s+(?:CONSTRAINT\s+{_IDENT}\s+)?"
+_ADD_FOREIGN_KEY: Final = re.compile(
+    rf"{_CONSTRAINT}FOREIGN\s+KEY\s*\(([^)]*)\)\s*REFERENCES\s+{_QUALIFIED}\s*\(([^)]*)\)(.*)$",
+    re.IGNORECASE | re.DOTALL,
+)
+_ADD_UNIQUE: Final = re.compile(rf"{_CONSTRAINT}UNIQUE\s*\(([^)]*)\)", re.IGNORECASE | re.DOTALL)
+_ADD_CHECK: Final = re.compile(rf"{_CONSTRAINT}CHECK\s*\((.*)\)\s*(NOT\s+VALID)?\s*$", re.IGNORECASE | re.DOTALL)
+_PLAIN_COLUMN: Final = re.compile(rf"^{_IDENT}(?:\s+(?:ASC|DESC|NULLS\s+(?:FIRST|LAST)))*$", re.IGNORECASE)
+_COMMENT: Final = re.compile(r"--[^\n]*")
+_NOT_NULL: Final = re.compile(r"\bNOT\s+NULL\b", re.IGNORECASE)
+_HAS_DEFAULT: Final = re.compile(r"\bDEFAULT\b|\bGENERATED\b", re.IGNORECASE)
+
+
+def migration_data_check_timeout_seconds() -> float:
+    """Seconds one count query may run before the preflight gives up on verifying it."""
+    raw: Final = os.getenv(MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR)
+    if raw is None:
+        return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
+    try:
+        return float(raw)
+    except ValueError:
+        logger.warning(
+            "%s=%r is not a number, using %s",
+            MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR,
+            raw,
+            DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS,
+        )
+        return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
+
+
+def statements(migration_sql: str) -> tuple[str, ...]:
+    """The statements of one migration file, comments stripped, empty ones dropped."""
+    parts: Final = _COMMENT.sub("", migration_sql).split(";")
+    return tuple(part.strip() for part in parts if part.strip())
+
+
+def _split_top_level(text: str) -> tuple[str, ...]:
+    """Split on the commas outside parentheses."""
+    depth: Final = accumulate((1 if char == "(" else -1 if char == ")" else 0) for char in text)
+    cuts: Final = tuple(
+        position for position, (char, level) in enumerate(zip(text, depth, strict=True)) if char == "," and level == 0
+    )
+    starts: Final = (0, *(cut + 1 for cut in cuts))
+    stops: Final = (*cuts, len(text))
+    pieces: Final = tuple(text[start:stop].strip() for start, stop in zip(starts, stops, strict=True))
+    return tuple(piece for piece in pieces if piece)
+
+
+def _plain_columns(column_list: str) -> tuple[str, ...] | None:
+    """The column names of an index or constraint column list, None when one entry is an
+    expression the derivation cannot turn into a count."""
+    matches: Final = tuple(_PLAIN_COLUMN.match(entry) for entry in _split_top_level(column_list))
+    if any(match is None for match in matches):
+        return None
+    return tuple(match.group(1) for match in matches if match is not None)
+
+
+def _index_checks(match: "re.Match[str]", statement_text: str, source: str) -> Iterator[Check]:
+    unique, name, table, method, column_list = match.groups()
+    columns: Final = _plain_columns(column_list)
+    if columns is None:
+        return
+    label: Final = f'CREATE {"UNIQUE " if unique else ""}INDEX "{name}"'
+    if method is None or method.lower() == "btree":
+        yield Check("too_long", table, columns, label, source, index=name)
+    if unique:
+        yield Check("duplicates", table, columns, label, source, index=name)
+
+
+def _alter_table_checks(table: str, action: str, source: str) -> Iterator[Check]:
+    if (set_not_null := _SET_NOT_NULL.match(action)) is not None:
+        yield Check("nulls", table, (set_not_null.group(1),), f'ALTER TABLE "{table}" {action}', source)
+        return
+    if (add_column := _ADD_COLUMN.match(action)) is not None:
+        column, body = add_column.groups()
+        if _NOT_NULL.search(body) is not None and _HAS_DEFAULT.search(body) is None:
+            yield Check(
+                "any_row", table, (column,), f'ALTER TABLE "{table}" ADD COLUMN "{column}" ... NOT NULL', source
+            )
+        return
+    if (foreign_key := _ADD_FOREIGN_KEY.match(action)) is not None:
+        constraint, columns, referenced_table, referenced_columns, tail = foreign_key.groups()
+        own: Final = _plain_columns(columns)
+        theirs: Final = _plain_columns(referenced_columns)
+        if own is None or theirs is None or len(own) != len(theirs) or "NOT VALID" in tail.upper():
+            return
+        yield Check(
+            "orphans",
+            table,
+            own,
+            _constraint_label(table, constraint, "FOREIGN KEY"),
+            source,
+            referenced_table=referenced_table,
+            referenced_columns=theirs,
+        )
+        return
+    if (unique := _ADD_UNIQUE.match(action)) is not None:
+        constraint, columns = unique.groups()
+        unique_columns: Final = _plain_columns(columns)
+        if unique_columns is not None:
+            yield Check("duplicates", table, unique_columns, _constraint_label(table, constraint, "UNIQUE"), source)
+        return
+    if (check := _ADD_CHECK.match(action)) is not None:
+        constraint, expression, not_valid = check.groups()
+        if not_valid is None:
+            yield Check(
+                "check_violations",
+                table,
+                (),
+                _constraint_label(table, constraint, "CHECK"),
+                source,
+                expression=expression,
+            )
+
+
+def _constraint_label(table: str, constraint: str | None, kind: str) -> str:
+    named: Final = f'CONSTRAINT "{constraint}" ' if constraint else ""
+    return f'ALTER TABLE "{table}" ADD {named}{kind}'
+
+
+def checks_for_statement(statement_text: str, source: str) -> tuple[Check, ...]:
+    """The count queries one SQL statement needs before it can be applied, empty for a
+    statement that cannot fail on existing rows or that the derivation cannot read."""
+    if (index := _CREATE_INDEX.match(statement_text)) is not None:
+        return tuple(_index_checks(index, statement_text, source))
+    if (alter := _ALTER_TABLE.match(statement_text)) is not None:
+        table, actions = alter.groups()
+        return tuple(
+            check for action in _split_top_level(actions) for check in _alter_table_checks(table, action, source)
+        )  # comprehension-ok: flattening the per-action checks of one statement
+    return ()
+
+
+def checks_for_migration(name: str, migration_sql: str) -> tuple[Check, ...]:
+    return tuple(
+        check
+        for statement_text in statements(migration_sql)
+        for check in checks_for_statement(statement_text, f"migration {name}")
+    )  # comprehension-ok: flattening the per-statement checks of one migration
+
+
+def checks_for_index(table: str, name: str, columns: tuple[str, ...]) -> tuple[Check, ...]:
+    """The entry-size check a request-log index needs before `CREATE INDEX CONCURRENTLY`."""
+    return (Check("too_long", table, columns, f'CREATE INDEX "{name}"', "request-log index", index=name),)
+
+
+def _applied_migrations(connection: "psycopg.Connection[tuple[object, ...]]", schema: str) -> frozenset[str]:
+    import psycopg
+    from psycopg import sql
+
+    try:
+        rows: Final = connection.execute(
+            sql.SQL("SELECT migration_name FROM {} WHERE finished_at IS NOT NULL AND rolled_back_at IS NULL").format(
+                sql.Identifier(schema, "_prisma_migrations")
+            )
+        ).fetchall()
+    except psycopg.errors.UndefinedTable:
+        return frozenset()
+    return frozenset(str(row[0]) for row in rows)
+
+
+def pending_migrations(
+    connection: "psycopg.Connection[tuple[object, ...]]", schema: str, migrations_dir: Path
+) -> tuple[tuple[str, str], ...]:
+    """(name, sql) of every shipped migration the ledger has not finished, oldest first.
+    Without a ledger every migration is pending, and the checks skip tables that do not exist yet."""
+    applied: Final = _applied_migrations(connection, schema)
+    folders: Final = tuple(folder for folder in sorted(migrations_dir.iterdir()) if folder.is_dir())
+    return tuple(
+        (folder.name, (folder / "migration.sql").read_text())
+        for folder in folders
+        if folder.name not in applied and (folder / "migration.sql").is_file()
+    )
+
+
+def _int(value: object, default: int) -> int:
+    return value if isinstance(value, int) else default
+
+
+def column_type_lookup(connection: "psycopg.Connection[tuple[object, ...]]", schema: str) -> ColumnTypeLookup:
+    def lookup(table: str, column: str) -> ColumnType | None:
+        row: Final = connection.execute(
+            "SELECT t.typname, a.attlen, a.attalign, a.atttypmod FROM pg_attribute a "
+            "JOIN pg_class c ON c.oid = a.attrelid JOIN pg_namespace n ON n.oid = c.relnamespace "
+            "JOIN pg_type t ON t.oid = a.atttypid "
+            "WHERE n.nspname = %s AND c.relname = %s AND a.attname = %s AND a.attnum > 0 AND NOT a.attisdropped",
+            (schema, table, column),
+        ).fetchone()
+        if row is None:
+            return None
+        return ColumnType(str(row[0]), _int(row[1], -1), str(row[2]), _int(row[3], -1))
+
+    return lookup
+
+
+def _column_types(check: Check, column_type: ColumnTypeLookup) -> tuple[ColumnType, ...] | None:
+    """The catalog types of the check's columns, None when one of them is not there yet."""
+    types: Final = tuple(column_type(check.table, column) for column in check.columns)
+    if any(column is None for column in types):
+        return None
+    return tuple(column for column in types if column is not None)
+
+
+def _table_exists(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> bool:
+    row: Final = connection.execute(
+        "SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name = %s", (schema, table)
+    ).fetchone()
+    return row is not None
+
+
+def _index_is_valid(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: str) -> bool:
+    row: Final = connection.execute(
+        "SELECT x.indisvalid FROM pg_index x JOIN pg_class i ON i.oid = x.indexrelid "
+        "JOIN pg_namespace n ON n.oid = i.relnamespace WHERE n.nspname = %s AND i.relname = %s",
+        (schema, index),
+    ).fetchone()
+    return row is not None and bool(row[0])
+
+
+def btree_max_entry_bytes(block_size: int) -> int:
+    """`BTMaxItemSize` from nbtree.h for btree version 4: a third of the page left after the page
+    header, three line pointers and the btree opaque area, rounded down to 8, minus the heap TID a
+    pivot tuple carries. 2704 on the default 8 KB page, the number in Postgres's error."""
+    usable: Final = block_size - _PAGE_HEADER_WITH_THREE_LINE_POINTERS_BYTES - _BTREE_PAGE_OPAQUE_BYTES
+    return (usable // 3) // 8 * 8 - _HEAP_TID_BYTES
+
+
+def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
+    from psycopg import sql
+
+    if column.type_name in _OCTET_LENGTH_TYPES:
+        return sql.SQL("octet_length({})").format(identifier)
+    return sql.SQL("octet_length({}::text)").format(identifier)
+
+
+def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":
+    from psycopg import sql
+
+    unit: Final = _ALIGNMENT.get(align, 8)
+    return sql.SQL("(({off} + {pad}) / {unit}) * {unit}").format(
+        off=offset, pad=sql.Literal(unit - 1), unit=sql.Literal(unit)
+    )
+
+
+def _value_end(column: ColumnType, identifier: "sql.Identifier", offset: "sql.Composable") -> "sql.Composed":
+    """Where this column's value ends when laid down from `offset` the way `index_form_tuple`
+    does it: NULL takes no room, a fixed-length value is aligned to its type, a variable-length
+    value of up to 126 bytes gets a 1-byte header and no alignment, a longer one a 4-byte header."""
+    from psycopg import sql
+
+    length: Final = _value_bytes(column, identifier)
+    placed: Final = (
+        sql.SQL("{} + {}").format(_aligned(offset, column.align), sql.Literal(column.length))
+        if not column.is_variable_length
+        else sql.SQL("CASE WHEN {len} <= {short} THEN {off} + 1 + {len} ELSE {aligned} + {hdr} + {len} END").format(
+            len=length,
+            short=sql.Literal(SHORT_VARLENA_MAX_BYTES),
+            off=offset,
+            aligned=_aligned(offset, column.align),
+            hdr=sql.Literal(VARLENA_HEADER_BYTES),
+        )
+    )
+    return sql.SQL("CASE WHEN {col} IS NULL THEN {off} ELSE {placed} END").format(
+        col=identifier, off=offset, placed=placed
+    )
+
+
+def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
+    """count(*) and the largest entry over the rows whose btree entry would not fit, plus the
+    longest value of each variable-length column: an 8-byte tuple header, 8 more for the null
+    bitmap when a column is NULL, each value laid down in index order, the total rounded up to 8."""
+    from psycopg import sql
+
+    identifiers: Final = tuple(sql.Identifier(column) for column in check.columns)
+    any_null: Final = sql.SQL(" OR ").join(sql.SQL("{} IS NULL").format(identifier) for identifier in identifiers)
+    header: Final = sql.SQL("{} + CASE WHEN {} THEN {} ELSE 0 END").format(
+        sql.Literal(INDEX_TUPLE_HEADER_BYTES), any_null, sql.Literal(NULL_BITMAP_BYTES)
+    )
+    innermost: Final = sql.SQL("SELECT {cols}, {hdr} AS off0 FROM {t}").format(
+        cols=sql.SQL(", ").join(identifiers), hdr=header, t=sql.Identifier(schema, check.table)
+    )
+
+    def lay_down(inner: "sql.Composed", step: tuple[int, tuple[ColumnType, "sql.Identifier"]]) -> "sql.Composed":
+        position, (column, identifier) = step
+        return sql.SQL("SELECT *, {end} AS {next} FROM ({inner}) AS {alias}").format(
+            end=_value_end(column, identifier, sql.Identifier(f"off{position}")),
+            next=sql.Identifier(f"off{position + 1}"),
+            inner=inner,
+            alias=sql.Identifier(f"s{position}"),
+        )
+
+    laid_out: Final = reduce(lay_down, enumerate(zip(types, identifiers, strict=True)), innermost)
+    entry: Final = _aligned(sql.Identifier(f"off{len(types)}"), "d")
+    longest: Final = sql.SQL("").join(
+        sql.SQL(", max({})").format(_value_bytes(column, identifier))
+        for column, identifier in zip(types, identifiers, strict=True)
+        if column.is_variable_length
+    )
+    return sql.SQL("SELECT count(*), max({entry}){longest} FROM ({rows}) AS entries WHERE {entry} > {limit}").format(
+        entry=entry, longest=longest, rows=laid_out, limit=sql.Literal(limit_bytes)
+    )
+
+
+def _count_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
+    from psycopg import sql
+
+    table: Final = sql.Identifier(schema, check.table)
+    columns: Final = sql.SQL(", ").join(sql.Identifier(column) for column in check.columns)
+    match check.kind:
+        case "too_long":
+            return _entry_size_query(check, schema, types, limit_bytes)
+        case "duplicates":
+            not_null: Final = sql.SQL(" AND ").join(
+                sql.SQL("{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
+            )
+            return sql.SQL(
+                "SELECT count(*) FROM (SELECT 1 FROM {t} WHERE {nn} GROUP BY {cols} HAVING count(*) > 1) AS d"
+            ).format(t=table, nn=not_null, cols=columns)
+        case "nulls":
+            return sql.SQL("SELECT count(*) FROM {t} WHERE {col} IS NULL").format(
+                t=table, col=sql.Identifier(check.columns[0])
+            )
+        case "any_row":
+            return sql.SQL("SELECT count(*) FROM {t}").format(t=table)
+        case "orphans":
+            pairs: Final = sql.SQL(" AND ").join(
+                sql.SQL("parent.{} = child.{}").format(sql.Identifier(theirs), sql.Identifier(own))
+                for own, theirs in zip(check.columns, check.referenced_columns, strict=True)
+            )
+            present: Final = sql.SQL(" AND ").join(
+                sql.SQL("child.{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
+            )
+            return sql.SQL(
+                "SELECT count(*) FROM {t} AS child WHERE {present} AND NOT EXISTS (SELECT 1 FROM {ref} AS parent WHERE {pairs})"
+            ).format(t=table, present=present, ref=sql.Identifier(schema, check.referenced_table or ""), pairs=pairs)
+        case "check_violations":
+            return sql.SQL("SELECT count(*) FROM {t} WHERE NOT ({expr})").format(
+                t=table,
+                expr=sql.SQL(check.expression or "true"),  # pyright: ignore[reportArgumentType]  # copied from our own migration file
+            )
+
+
+def _applies(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    check: Check,
+    types: tuple[ColumnType, ...] | None,
+) -> bool:
+    """False when the check has nothing to count yet: the table or a column is not there
+    (a pending migration adds it), the index is already valid, or no column can overflow."""
+    if not _table_exists(connection, schema, check.table):
+        return False
+    if check.index is not None and _index_is_valid(connection, schema, check.index):
+        return False
+    if check.kind == "any_row":
+        return types is None
+    if types is None:
+        return False
+    if check.kind == "too_long":
+        return any(column.can_overflow_btree for column in types)
+    if check.kind == "orphans":
+        return check.referenced_table is not None and _table_exists(connection, schema, check.referenced_table)
+    return True
+
+
+def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str:
+    variable: Final = tuple(
+        column for column, column_type in zip(check.columns, types, strict=True) if column_type.is_variable_length
+    )
+    longest: Final = ", ".join(
+        f"{column} up to {length} bytes"
+        for column, length in zip(variable, values[2:], strict=True)
+        if length is not None
+    )
+    return f"with an index entry over {limit_bytes} bytes, largest {values[1]} bytes ({longest})"
+
+
+def _run_check(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    check: Check,
+    types: tuple[ColumnType, ...],
+    limit_bytes: int,
+    timeout_seconds: float,
+) -> Finding | None:
+    import psycopg
+
+    try:
+        row: Final = connection.execute(_count_query(check, schema, types, limit_bytes)).fetchone()
+    except psycopg.errors.QueryCanceled:
+        return CouldNotVerify(check, timeout_seconds)
+    values: Final[tuple[object, ...]] = row if row is not None else ()
+    rows: Final = _int(values[0], 0) if values else 0
+    if rows == 0:
+        return None
+    match check.kind:
+        case "too_long":
+            return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
+        case "duplicates":
+            return BlockedChange(check, rows, "duplicated value groups")
+        case "nulls":
+            return BlockedChange(check, rows, "NULL")
+        case "any_row":
+            return BlockedChange(
+                check, rows, "rows, and a NOT NULL column with no DEFAULT cannot be added to a table with rows"
+            )
+        case "orphans":
+            return BlockedChange(check, rows, f'with no matching row in "{check.referenced_table}"')
+        case "check_violations":
+            return BlockedChange(check, rows, f"failing CHECK ({check.expression})")
+
+
+def find_blockers(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    checks: tuple[Check, ...],
+    timeout_seconds: float = DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS,
+) -> tuple[Finding, ...]:
+    """Run every applicable check read-only, each under `timeout_seconds`. Never writes
+    and never takes a table lock; a check that times out is a finding, not a pass."""
+    from psycopg import sql
+
+    previous_timeout: Final = connection.execute("SHOW statement_timeout").fetchone()
+    previous_read_only: Final = connection.execute("SHOW default_transaction_read_only").fetchone()
+    block_size: Final = connection.execute("SHOW block_size").fetchone()
+    limit_bytes: Final = btree_max_entry_bytes(int(str(block_size[0])) if block_size else DEFAULT_BLOCK_SIZE)
+    connection.execute("SET default_transaction_read_only = on")
+    connection.execute(sql.SQL("SET statement_timeout = {}").format(sql.Literal(int(timeout_seconds * 1000))))
+    column_type: Final = column_type_lookup(connection, schema)
+    typed: Final = tuple((check, _column_types(check, column_type)) for check in checks)
+    applicable: Final = tuple(
+        (check, types if types is not None else ())
+        for check, types in typed
+        if _applies(connection, schema, check, types)
+    )
+    results: Final = tuple(
+        _run_check(connection, schema, check, types, limit_bytes, timeout_seconds) for check, types in applicable
+    )
+    connection.execute(
+        sql.SQL("SET statement_timeout = {}").format(sql.Literal(str(previous_timeout[0]) if previous_timeout else "0"))
+    )
+    connection.execute(
+        sql.SQL("SET default_transaction_read_only = {}").format(
+            sql.Literal(str(previous_read_only[0]) if previous_read_only else "off")
+        )
+    )
+    return tuple(result for result in results if result is not None)
+
+
+def _remedy(check: Check) -> str:
+    if check.kind == "too_long" and check.table == "LiteLLM_SpendLogs" and "api_key" in check.columns:
+        return LEGACY_SPEND_LOG_KEY_REMEDY
+    joined: Final = ", ".join(check.columns)
+    match check.kind:
+        case "too_long":
+            return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
+        case "duplicates":
+            return f'Remove or merge the duplicate "{check.table}" ({joined}) values.'
+        case "nulls":
+            return f'Fill the NULL "{check.table}"."{joined}" values, or give the column a default first.'
+        case "any_row":
+            return "Give the new column a DEFAULT in the migration, or add it nullable and fill it afterwards."
+        case "orphans":
+            return f'Delete or repoint the "{check.table}" ({joined}) values with no "{check.referenced_table}" row.'
+        case "check_violations":
+            return f'Fix the "{check.table}" rows that fail the expression before adding the constraint.'
+
+
+def _describe(finding: Finding) -> str:
+    check: Final = finding.check
+    where: Final = f'"{check.table}"' + (f".{', '.join(check.columns)}" if check.columns else "")
+    match finding:
+        case CouldNotVerify(timeout_seconds=seconds):
+            return (
+                f"  {where}: could not be checked within {seconds:g} s, refusing rather than guessing\n"
+                f"    blocks: {check.statement} ({check.source})\n"
+                f"    raise {MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR} to give the scan more time"
+            )
+        case BlockedChange(rows=rows, detail=detail):
+            return (
+                f"  {where}: {rows} row(s) {detail}\n"
+                f"    blocks: {check.statement} ({check.source})\n"
+                f"    fix:    {_remedy(check).replace(chr(10), chr(10) + ' ' * 8)}"
+            )
+
+
+def report(findings: tuple[Finding, ...]) -> str:
+    """The operator-facing refusal: one block per finding, then what to do next."""
+    lines: Final = "\n".join(_describe(finding) for finding in findings)
+    return (
+        f"{REFUSAL}\n{lines}\n"
+        "No schema change was made. Fix the rows and start again, or run `litellm --check_migrations_against_existing_data` with the new image to check again."
+    )
+
+
+def pending_change_checks(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    migrations_dir: Path | None,
+    indexes: tuple[tuple[str, str, tuple[str, ...]], ...],
+) -> tuple[Check, ...]:
+    """Every check the pending migrations and the request-log indexes need. `migrations_dir`
+    is None when the schema is pushed rather than migrated, so only the indexes are checked."""
+    migration_checks: Final = (
+        tuple(
+            check
+            for name, migration_sql in pending_migrations(connection, schema, migrations_dir)
+            for check in checks_for_migration(name, migration_sql)
+        )  # comprehension-ok: flattening the per-migration checks
+        if migrations_dir is not None
+        else ()
+    )
+    index_checks: Final = tuple(
+        check for table, name, columns in indexes for check in checks_for_index(table, name, columns)
+    )  # comprehension-ok: flattening the per-index checks
+    return migration_checks + index_checks

litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py

@@ -18,6 +18,12 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Final
 
 from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras.migration_data_checks import (
+    checks_for_index,
+    find_blockers,
+    migration_data_check_timeout_seconds,
+    report,
+)
 from litellm_proxy_extras.migration_lock import held_migration_lock
 
 if TYPE_CHECKING:
@@ -355,6 +361,12 @@ def _build_leaf_index(
             connection.execute(sql.SQL("DROP INDEX CONCURRENTLY {}").format(sql.Identifier(schema, name)))
         elif _adopt_equivalent_index(connection, schema, table, name, index):
             return True
+        blockers: Final = find_blockers(
+            connection, schema, checks_for_index(table, name, index.columns), migration_data_check_timeout_seconds()
+        )
+        if blockers:
+            logger.warning("Not building index %s on %s. %s", name, table, report(blockers))
+            return False
         logger.info("Building index %s on %s concurrently", name, table)
         prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format(
             sql.Identifier(name), sql.Identifier(schema, table)

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -16,6 +16,13 @@ from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
 from litellm_proxy_extras import prisma_toolchain
 from litellm_proxy_extras._logging import logger
 from litellm_proxy_extras._redaction import _redact_credentials
+from litellm_proxy_extras.migration_data_checks import (
+    Finding,
+    find_blockers,
+    migration_data_check_timeout_seconds,
+    pending_change_checks,
+    report,
+)
 from litellm_proxy_extras.migration_lock import held_migration_lock
 from litellm_proxy_extras.prisma_toolchain import (
     PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@@ -28,7 +35,11 @@ from litellm_proxy_extras.replica_identity import (
     REPLICA_IDENTITY_FULL_ENV_VAR,
     apply_replica_identity_full,
 )
-from litellm_proxy_extras.request_log_indexes import ensure_request_log_indexes, filter_request_log_index_diff
+from litellm_proxy_extras.request_log_indexes import (
+    REQUEST_LOG_INDEXES,
+    ensure_request_log_indexes,
+    filter_request_log_index_diff,
+)
 
 if TYPE_CHECKING:
     import psycopg
@@ -53,9 +64,7 @@ def _get_prisma_env() -> dict:
     if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
         # These env vars prevent Prisma from attempting downloads
         prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true"
-        prisma_env["NPM_CONFIG_CACHE"] = os.getenv(
-            "NPM_CONFIG_CACHE", "/app/.cache/npm"
-        )
+        prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm")
     return prisma_env
 
 
@@ -86,6 +95,7 @@ class _InvalidIndex:
     name: str
     table_size: str
 
+
 MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
 LIBPQ_URL_PARAMS: Final = frozenset(
     {
@@ -135,9 +145,7 @@ class _MigrateAttemptBudget:
 
 
 _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
-_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
-    r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
-)
+_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE)
 _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
     r'^(?:DROP\s+CONSTRAINT\s+"[^"]*_pkey"'
     r'|ADD\s+(?:CONSTRAINT\s+"[^"]*"\s+)?PRIMARY\s+KEY\s*\([^)]*\))$',
@@ -146,9 +154,9 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
 
 PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
     "LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
-    "so its primary key must include the partition key (\"startTime\"). `prisma db push` "
+    'so its primary key must include the partition key ("startTime"). `prisma db push` '
     "reconciles the database against schema.prisma, which declares the unpartitioned "
-    "primary key (\"request_id\"), and Postgres rejects that rewrite with: unique "
+    'primary key ("request_id"), and Postgres rejects that rewrite with: unique '
     "constraint on partitioned table must include all partitioning columns. Start the "
     "proxy without --use_prisma_db_push so it uses `prisma migrate deploy`, which only "
     "applies shipped migrations and leaves the partitioned primary key alone."
@@ -157,9 +165,7 @@ PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
 
 def _without_sql_comments(statement: str) -> str:
     return "\n".join(
-        line
-        for line in statement.splitlines()
-        if line.strip() and not line.strip().startswith("--")
+        line for line in statement.splitlines() if line.strip() and not line.strip().startswith("--")
     ).strip()
 
 
@@ -169,7 +175,7 @@ def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
         return statement
     kept = tuple(
         clause.strip()
-        for clause in statement[prefix_match.end():].split(",\n")
+        for clause in statement[prefix_match.end() :].split(",\n")
         if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause.strip())
     )
     if not kept:
@@ -243,8 +249,7 @@ def _get_prisma_command() -> str:
 
         # If not found, log warning and fall back
         logger.warning(
-            f"Prisma CLI not found at {default_cli_path}. "
-            "Falling back to Python wrapper (may attempt downloads)"
+            f"Prisma CLI not found at {default_cli_path}. Falling back to Python wrapper (may attempt downloads)"
         )
 
     # Fall back to the Python wrapper (will work in online mode)
@@ -331,9 +336,7 @@ class ProxyExtrasDBManager:
 
             return True
         except subprocess.TimeoutExpired:
-            logger.warning(
-                "Migration timed out - the database might be under heavy load."
-            )
+            logger.warning("Migration timed out - the database might be under heavy load.")
             return False
         except subprocess.CalledProcessError as e:
             logger.warning(
@@ -386,10 +389,12 @@ class ProxyExtrasDBManager:
         except ImportError:
             return None
 
+        from psycopg import sql
+
         cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
-        ledger_table: Final = psycopg.sql.SQL("{}.{}").format(
-            psycopg.sql.Identifier(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"),
-            psycopg.sql.Identifier("_prisma_migrations"),
+        ledger_table: Final = sql.SQL("{}.{}").format(
+            sql.Identifier(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"),
+            sql.Identifier("_prisma_migrations"),
         )
         try:
             with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
@@ -500,9 +505,7 @@ class ProxyExtrasDBManager:
         return filter_partitioned_spend_logs_diff(without_indexes)
 
     @staticmethod
-    def _resolve_all_migrations(
-        migrations_dir: str, schema_path: str, mark_all_applied: bool = True
-    ):
+    def _resolve_all_migrations(migrations_dir: str, schema_path: str, mark_all_applied: bool = True):
         """
         1. Compare the current database state to schema.prisma and generate a migration for the diff.
         2. Run prisma migrate deploy to apply any pending migrations.
@@ -571,9 +574,7 @@ class ProxyExtrasDBManager:
                     )
                     logger.info(f"Applied migration: {mig_file.parent.name}")
                 except subprocess.CalledProcessError as e:
-                    logger.warning(
-                        f"Failed to apply migration {mig_file.parent.name}: {e.stderr}"
-                    )
+                    logger.warning(f"Failed to apply migration {mig_file.parent.name}: {e.stderr}")
                 except subprocess.TimeoutExpired:
                     logger.warning(f"Migration {mig_file.parent.name} timed out.")
             return
@@ -645,9 +646,7 @@ class ProxyExtrasDBManager:
                 logger.debug(f"Resolved migration: {migration_name}")
             except subprocess.CalledProcessError as e:
                 if "is already recorded as applied in the database." not in e.stderr:
-                    logger.warning(
-                        f"Failed to resolve migration {migration_name}: {e.stderr}"
-                    )
+                    logger.warning(f"Failed to resolve migration {migration_name}: {e.stderr}")
 
     @staticmethod
     def raise_if_lens_rename_pending(connect: LensCheckConnect | None = None) -> None:
@@ -707,9 +706,7 @@ class ProxyExtrasDBManager:
 
         cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
         try:
-            with psycopg.connect(
-                cleaned_url, connect_timeout=10, autocommit=True
-            ) as conn:
+            with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
                 row = conn.execute(
                     "SELECT 1 "
                     "FROM pg_partitioned_table pt "
@@ -717,10 +714,7 @@ class ProxyExtrasDBManager:
                     "JOIN pg_namespace n ON n.oid = c.relnamespace "
                     "WHERE c.relname = 'LiteLLM_SpendLogs' "
                     "  AND n.nspname = %s",
-                    (
-                        ProxyExtrasDBManager._prisma_schema_param(database_url)
-                        or "public",
-                    ),
+                    (ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",),
                 ).fetchone()
         except (psycopg.OperationalError, psycopg.DatabaseError):
             return False
@@ -815,9 +809,7 @@ class ProxyExtrasDBManager:
             # it, psycopg3's `with conn` calls COMMIT on clean exit — which
             # fails after `UndefinedTable` (fresh DB) leaves the transaction
             # in an aborted state.
-            with psycopg.connect(
-                cleaned_url, connect_timeout=10, autocommit=True
-            ) as conn:
+            with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
                 try:
                     rows = conn.execute(
                         "SELECT migration_name FROM _prisma_migrations "
@@ -838,9 +830,7 @@ class ProxyExtrasDBManager:
             return
 
         head_newest_ts = _max_migration_timestamp(known)
-        hostile = {
-            name for name in unknown if _migration_timestamp(name) > head_newest_ts
-        }
+        hostile = {name for name in unknown if _migration_timestamp(name) > head_newest_ts}
         if not hostile:
             return
 
@@ -931,8 +921,7 @@ class ProxyExtrasDBManager:
                 if not found:
                     return True
                 logger.warning(
-                    "Found %d invalid index(es) left by an interrupted CREATE INDEX "
-                    "CONCURRENTLY, rebuilding: %s",
+                    "Found %d invalid index(es) left by an interrupted CREATE INDEX CONCURRENTLY, rebuilding: %s",
                     len(found),
                     ", ".join(f"{index.name} (table size {index.table_size})" for index in found),
                 )
@@ -1403,8 +1392,7 @@ class ProxyExtrasDBManager:
             prisma_env = _get_prisma_env()
         except OSError as e:
             logger.error(
-                "Could not resolve the migrations directory for the REPLICA "
-                "IDENTITY FULL step, skipping it. Error: %s",
+                "Could not resolve the migrations directory for the REPLICA IDENTITY FULL step, skipping it. Error: %s",
                 e,
             )
             return False
@@ -1435,18 +1423,57 @@ class ProxyExtrasDBManager:
         Returns:
             bool: True if setup was successful, False otherwise
         """
-        ensure_prisma_toolchain(
-            prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env()
-        )
-        migrated = ProxyExtrasDBManager._run_migrations(
-            use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
-        )
+        ensure_prisma_toolchain(prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env())
+        ProxyExtrasDBManager.raise_if_rows_block_pending_changes(use_migrate=use_migrate)
+        migrated = ProxyExtrasDBManager._run_migrations(use_migrate=use_migrate, use_v2_resolver=use_v2_resolver)
         if not migrated:
             return False
         ProxyExtrasDBManager.repair_invalid_indexes()
         ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
         return True
 
+    @staticmethod
+    def migration_data_check_findings(use_migrate: bool = True) -> tuple[Finding, ...]:
+        """What in the existing rows would make a pending migration or request-log index
+        fail, read with count queries only: no write, no table lock. Empty when nothing
+        blocks, when no database is configured, or when it cannot be reached, so the
+        migration step that follows reports the connection problem the way it does today."""
+        import psycopg
+
+        database_url: Final = os.environ.get("DATABASE_URL")
+        if not database_url:
+            return ()
+        direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
+            os.environ.get("DIRECT_URL") or database_url
+        )
+        schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
+        migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" if use_migrate else None
+        indexes: Final = tuple((index.table, index.name, index.columns) for index in REQUEST_LOG_INDEXES)
+        try:
+            with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection:
+                checks: Final = pending_change_checks(connection, schema, migrations_dir, indexes)
+                return find_blockers(connection, schema, checks, migration_data_check_timeout_seconds())
+        except psycopg.OperationalError as exc:
+            logger.warning(
+                "Could not reach the database for the pre-migration row check: %s", _redact_credentials(str(exc))
+            )
+            return ()
+
+    @staticmethod
+    def migration_data_check_report(use_migrate: bool = True) -> Optional[str]:
+        """The operator-facing refusal for `migration_data_check_findings`, None when nothing blocks."""
+        findings: Final = ProxyExtrasDBManager.migration_data_check_findings(use_migrate=use_migrate)
+        return report(findings) if findings else None
+
+    @staticmethod
+    def raise_if_rows_block_pending_changes(use_migrate: bool) -> None:
+        """Refuse to run migrations that the existing rows would make fail, before any is
+        applied. The RuntimeError carries the report and reaches the operator through the
+        proxy's "Database migration cannot proceed" exit."""
+        refusal: Final = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate)
+        if refusal is not None:
+            raise RuntimeError(refusal)
+
     @staticmethod
     def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
         """Build the indexes in `REQUEST_LOG_INDEXES` on the writer, in the schema the
@@ -1513,9 +1540,7 @@ class ProxyExtrasDBManager:
                         # Skip sanity check when deploy reports no pending migrations —
                         # DB already matches schema, no drift to correct.
                         if "No pending migrations to apply" in result.stdout:
-                            logger.info(
-                                "No pending migrations — skipping post-migration sanity check"
-                            )
+                            logger.info("No pending migrations — skipping post-migration sanity check")
                             return True
 
                         # Run sanity check to ensure DB matches schema
@@ -1530,9 +1555,7 @@ class ProxyExtrasDBManager:
                         logger.info(f"prisma db error: {stderr}, e: {e.stdout}")
                         if "P3009" in stderr:
                             # Extract the failed migration name from the error message
-                            migration_match = re.search(
-                                r"`(\d+_.*)` migration", stderr
-                            )
+                            migration_match = re.search(r"`(\d+_.*)` migration", stderr)
                             if migration_match:
                                 failed_migration = migration_match.group(1)
                                 if ProxyExtrasDBManager._is_idempotent_error(stderr):
@@ -1540,9 +1563,7 @@ class ProxyExtrasDBManager:
                                         f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied"
                                     )
                                     try:
-                                        ProxyExtrasDBManager._roll_back_migration(
-                                            failed_migration
-                                        )
+                                        ProxyExtrasDBManager._roll_back_migration(failed_migration)
                                     except (
                                         subprocess.CalledProcessError,
                                         subprocess.TimeoutExpired,
@@ -1552,9 +1573,7 @@ class ProxyExtrasDBManager:
                                             f"It may already be in a rolled-back state."
                                         )
                                     try:
-                                        ProxyExtrasDBManager._resolve_specific_migration(
-                                            failed_migration
-                                        )
+                                        ProxyExtrasDBManager._resolve_specific_migration(failed_migration)
                                         logger.info(
                                             f"✅ Migration {failed_migration} resolved, retrying to apply remaining migrations"
                                         )
@@ -1562,9 +1581,7 @@ class ProxyExtrasDBManager:
                                         subprocess.CalledProcessError,
                                         subprocess.TimeoutExpired,
                                     ) as resolve_err:
-                                        logger.warning(
-                                            f"Failed to resolve migration {failed_migration}: {resolve_err}"
-                                        )
+                                        logger.warning(f"Failed to resolve migration {failed_migration}: {resolve_err}")
                                     # Apply any schema drift not covered by the marked-as-applied migration
                                     ProxyExtrasDBManager._resolve_all_migrations(
                                         migrations_dir,
@@ -1572,9 +1589,7 @@ class ProxyExtrasDBManager:
                                         mark_all_applied=False,
                                     )
                                 else:
-                                    logger.info(
-                                        f"Found failed migration: {failed_migration}, marking as rolled back"
-                                    )
+                                    logger.info(f"Found failed migration: {failed_migration}, marking as rolled back")
                                     # Mark the failed migration as rolled back
                                     prisma_toolchain.run_prisma(
                                         [
@@ -1587,23 +1602,14 @@ class ProxyExtrasDBManager:
                                         timeout=prisma_command_timeout(),
                                         env=_get_prisma_env(),
                                     )
-                                    logger.info(
-                                        f"✅ Migration {failed_migration} marked as rolled back... retrying"
-                                    )
-                        elif (
-                            "P3005" in stderr
-                            and "database schema is not empty" in stderr
-                        ):
+                                    logger.info(f"✅ Migration {failed_migration} marked as rolled back... retrying")
+                        elif "P3005" in stderr and "database schema is not empty" in stderr:
                             logger.info(
                                 "Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system"
                             )
                             ProxyExtrasDBManager._create_baseline_migration(schema_path)
-                            logger.info(
-                                "Baseline migration created, resolving all migrations"
-                            )
-                            ProxyExtrasDBManager._resolve_all_migrations(
-                                migrations_dir, schema_path
-                            )
+                            logger.info("Baseline migration created, resolving all migrations")
+                            ProxyExtrasDBManager._resolve_all_migrations(migrations_dir, schema_path)
                             logger.info("✅ All migrations resolved.")
                             return True
                         elif "P3018" in stderr:
@@ -1611,14 +1617,8 @@ class ProxyExtrasDBManager:
                             if ProxyExtrasDBManager._is_permission_error(stderr):
                                 # Permission errors should NOT be marked as applied
                                 # Extract migration name for logging
-                                migration_match = re.search(
-                                    r"Migration name: (\d+_.*)", stderr
-                                )
-                                migration_name = (
-                                    migration_match.group(1)
-                                    if migration_match
-                                    else "unknown"
-                                )
+                                migration_match = re.search(r"Migration name: (\d+_.*)", stderr)
+                                migration_name = migration_match.group(1) if migration_match else "unknown"
 
                                 logger.error(
                                     f"❌ Migration {migration_name} failed due to insufficient permissions. "
@@ -1628,16 +1628,10 @@ class ProxyExtrasDBManager:
                                 # Mark as rolled back and exit with error
                                 if migration_match:
                                     try:
-                                        ProxyExtrasDBManager._roll_back_migration(
-                                            migration_name
-                                        )
-                                        logger.info(
-                                            f"Migration {migration_name} marked as rolled back"
-                                        )
+                                        ProxyExtrasDBManager._roll_back_migration(migration_name)
+                                        logger.info(f"Migration {migration_name} marked as rolled back")
                                     except Exception as rollback_error:
-                                        logger.warning(
-                                            f"Failed to mark migration as rolled back: {rollback_error}"
-                                        )
+                                        logger.warning(f"Failed to mark migration as rolled back: {rollback_error}")
 
                                 # Re-raise the error to prevent silent failures
                                 raise RuntimeError(
@@ -1652,18 +1646,12 @@ class ProxyExtrasDBManager:
                                     "resolving as applied"
                                 )
                                 # Extract the migration name from the error message
-                                migration_match = re.search(
-                                    r"Migration name: (\d+_.*)", stderr
-                                )
+                                migration_match = re.search(r"Migration name: (\d+_.*)", stderr)
                                 if migration_match:
                                     migration_name = migration_match.group(1)
                                     try:
-                                        logger.info(
-                                            f"Rolling back migration {migration_name}"
-                                        )
-                                        ProxyExtrasDBManager._roll_back_migration(
-                                            migration_name
-                                        )
+                                        logger.info(f"Rolling back migration {migration_name}")
+                                        ProxyExtrasDBManager._roll_back_migration(migration_name)
                                     except (
                                         subprocess.CalledProcessError,
                                         subprocess.TimeoutExpired,
@@ -1677,9 +1665,7 @@ class ProxyExtrasDBManager:
                                             f"Resolving migration {migration_name} that failed "
                                             f"due to existing schema objects"
                                         )
-                                        ProxyExtrasDBManager._resolve_specific_migration(
-                                            migration_name
-                                        )
+                                        ProxyExtrasDBManager._resolve_specific_migration(migration_name)
                                         logger.info(
                                             f"✅ Migration {migration_name} resolved, "
                                             f"retrying to apply remaining migrations"
@@ -1688,9 +1674,7 @@ class ProxyExtrasDBManager:
                                         subprocess.CalledProcessError,
                                         subprocess.TimeoutExpired,
                                     ) as resolve_err:
-                                        logger.warning(
-                                            f"Failed to resolve migration {migration_name}: {resolve_err}"
-                                        )
+                                        logger.warning(f"Failed to resolve migration {migration_name}: {resolve_err}")
                                     # Apply any schema drift not covered by the marked-as-applied migration
                                     ProxyExtrasDBManager._resolve_all_migrations(
                                         migrations_dir,
@@ -1732,14 +1716,8 @@ class ProxyExtrasDBManager:
                 time.sleep(random.randrange(5, 15))
             except subprocess.CalledProcessError as e:
                 attempts_left = 3 - attempt
-                retry_msg = (
-                    f" Retrying... ({attempts_left} attempts left)"
-                    if attempts_left > 0
-                    else ""
-                )
-                stderr_detail: Final = (
-                    f" stderr: {_redact_credentials(str(e.stderr))}" if e.stderr else ""
-                )
+                retry_msg = f" Retrying... ({attempts_left} attempts left)" if attempts_left > 0 else ""
+                stderr_detail: Final = f" stderr: {_redact_credentials(str(e.stderr))}" if e.stderr else ""
                 logger.error(
                     f"The process failed to execute. Details: {_redact_command_error(e)}.{stderr_detail}{retry_msg}"
                 )

litellm/litellm_core_utils/litellm_logging.py

@@ -97,8 +97,13 @@ from litellm.litellm_core_utils.internal_call_metadata import (
     is_unbilled_non_inference_call,
 )
 from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
+    bedrock_guardrail_trace,
     cost_breakdown_with_guardrail,
     guardrail_information_cost,
+    is_str_keyed_dict,
+    native_bedrock_guardrail_information,
+    object_mapping,
+    with_guardrail_entry,
 )
 from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
     StandardBuiltInToolCostTracking,
@@ -2725,6 +2730,31 @@ class Logging(LiteLLMLoggingBaseClass):
             return
         self.model_call_details["response_headers"] = headers
 
+    def record_native_bedrock_guardrail(self, result: object) -> None:
+        """Bill the guardrail trace a Bedrock response carries when the deployment's ``guardrailConfig`` ran it, by
+        recording it as guardrail information on this call, the same entry the proxy-side Bedrock guardrail writes.
+        A cache hit carries the original response's trace but made no Bedrock call, so it bills nothing."""
+        trace: Final = bedrock_guardrail_trace(result)
+        if trace is None or served_from_cache(result) or self.model_call_details.get("cache_hit") is True:
+            return
+        optional_params: Final = object_mapping(getattr(self, "optional_params", None)) or {}
+        guardrail_config: Final = object_mapping(optional_params.get("guardrailConfig")) or {}
+        identifier: Final = guardrail_config.get("guardrailIdentifier")
+        litellm_params_in: Final = object_mapping(getattr(self, "litellm_params", None)) or {}
+        region: Final = optional_params.get("aws_region_name") or litellm_params_in.get("aws_region_name")
+        entry: Final = native_bedrock_guardrail_information(
+            trace,
+            aws_region_name=region if isinstance(region, str) else None,
+            guardrail_identifier=identifier if isinstance(identifier, str) else None,
+        )
+        litellm_params: Final[object] = self.model_call_details.get("litellm_params")
+        if entry is None or not is_str_keyed_dict(litellm_params):
+            return
+        metadata: Final[object] = litellm_params.get("metadata")
+        bucket: Final = metadata if is_str_keyed_dict(metadata) else {}
+        litellm_params["metadata"] = bucket
+        with_guardrail_entry(bucket, entry)
+
     def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None:
         """
         Copy response._hidden_params into litellm_params.metadata['hidden_params'].
@@ -7111,12 +7141,12 @@ class StandardLoggingPayloadSetup:
         request_headers: Final = cast(  # cast-ok: request headers are untyped framework data
             dict[str, str], headers
         )
-
-        header_tags: Final = []
-        for header_name in extra_headers:
-            header_value = request_headers.get(header_name)
-            if header_value:
-                header_tags.append(f"{header_name}: {header_value}")
+        lowercase_request_headers: Final = {str(name).lower(): value for name, value in request_headers.items()}
+        header_tags: Final = [
+            f"{header_name}: {header_value}"
+            for header_name in extra_headers
+            if (header_value := lowercase_request_headers.get(header_name.lower()))
+        ]
 
         return header_tags if header_tags else None
 

litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py

@@ -1,15 +1,24 @@
 import math
-from collections.abc import Mapping
+from collections.abc import Iterator, Mapping, Sequence
+from functools import partial
+from itertools import chain
 from typing import Annotated, Final
 
-from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
+from typing_extensions import TypeIs  # noqa: TID251  # TypeIs reaches typing only on 3.13
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.core_helpers import (
+    redact_nested_match_and_regex_keys,  # pyright: ignore[reportUnknownVariableType]  # signature uses bare dict
+)
+from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.llms.base import LiteLLMBaseModel
-from litellm.types.utils import CostBreakdown
+from litellm.types.utils import CostBreakdown, GuardrailStatus, StandardLoggingGuardrailInformation
 
 BEDROCK_GUARDRAIL_PRICING_KEY: Final = "bedrock/guardrails"
+BEDROCK_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace")
+_INTERVENING_ACTIONS: Final = frozenset({"BLOCKED", "ANONYMIZED"})
 
 
 class GuardrailPricing(LiteLLMBaseModel):
@@ -96,6 +105,161 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str
     return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name))
 
 
+_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object])
+_OBJECT_ITEMS: Final = TypeAdapter(tuple[object, ...])
+
+
+def object_mapping(value: object) -> Mapping[str, object] | None:
+    """``value`` as a str-keyed mapping, None for anything else. Vendor JSON arrives untyped, so this is the one place
+    that turns it into something the rest of the module can read without guessing."""
+    try:
+        return _OBJECT_MAPPING.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def object_items(value: object) -> tuple[object, ...]:
+    if isinstance(value, (str, bytes, Mapping)):
+        return ()
+    try:
+        return _OBJECT_ITEMS.validate_python(value)
+    except ValidationError:
+        return ()
+
+
+def is_str_keyed_dict(
+    value: object,
+) -> TypeIs[dict[str, object]]:  # guard-ok: JSON object keys are str  # mutable-ok: the caller edits this dict in place
+    return isinstance(value, dict)
+
+
+def _trace_field(response: object, field: str) -> Mapping[str, object] | None:
+    source: Final = object_mapping(response)
+    trace: Final = object_mapping(source.get(field) if source is not None else getattr(response, field, None))
+    if trace is None or object_mapping(trace.get("guardrail")) is None:
+        return None
+    return trace
+
+
+def bedrock_guardrail_trace(response: object) -> Mapping[str, object] | None:
+    """The guardrail trace a Bedrock response carries: Converse puts it under ``trace``, Invoke under
+    ``amazon-bedrock-trace``. Returned as is, or None when the response has no guardrail section."""
+    traces: Final = map(partial(_trace_field, response), BEDROCK_TRACE_FIELDS)
+    return next((trace for trace in traces if trace is not None), None)
+
+
+class _Assessment(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="allow")
+    invocation_metrics: Mapping[str, JsonValue] = Field(default_factory=dict, alias="invocationMetrics")
+
+    @property
+    def billable_usage(self) -> Mapping[str, int]:
+        usage: Final = self.invocation_metrics.get("usage")
+        if not isinstance(usage, dict):
+            return {}
+        return {
+            counter: units
+            for counter, units in usage.items()
+            if counter.endswith("Units") and isinstance(units, int) and not isinstance(units, bool)
+        }
+
+
+class _GuardrailTrace(BaseModel):
+    """Both trace shapes: Converse ``inputAssessment`` {id: a} and ``outputAssessments`` {id: [a, ...]}, Invoke
+    ``input`` {id: a} and ``outputs`` [{id: a}, ...]."""
+
+    model_config = ConfigDict(frozen=True, extra="allow")
+    input_assessment: Mapping[str, _Assessment] = Field(default_factory=dict, alias="inputAssessment")
+    input: Mapping[str, _Assessment] = Field(default_factory=dict)
+    output_assessments: Mapping[str, tuple[_Assessment, ...]] = Field(default_factory=dict, alias="outputAssessments")
+    outputs: tuple[Mapping[str, _Assessment], ...] = ()
+
+    @property
+    def assessments(self) -> tuple[tuple[str, _Assessment], ...]:
+        listed: Final = (
+            (gid, a) for gid, items in self.output_assessments.items() for a in items
+        )  # comprehension-ok: id then its list
+        return (
+            *self.input_assessment.items(),
+            *self.input.items(),
+            *listed,
+            *chain.from_iterable(by_id.items() for by_id in self.outputs),
+        )
+
+
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+
+def _summed_usage(usages: Sequence[Mapping[str, int]]) -> Mapping[str, int]:
+    counters: Final = tuple(dict.fromkeys(chain.from_iterable(usages)))
+    return {counter: sum(usage.get(counter, 0) for usage in usages) for counter in counters}
+
+
+def _actions(node: JsonValue) -> Iterator[str]:
+    if isinstance(node, list):
+        yield from chain.from_iterable(map(_actions, node))
+        return
+    if not isinstance(node, dict):
+        return
+    action: Final = node.get("action")
+    if isinstance(action, str):
+        yield action
+    yield from chain.from_iterable(map(_actions, node.values()))
+
+
+def native_bedrock_guardrail_information(
+    trace: Mapping[str, object], aws_region_name: str | None, guardrail_identifier: str | None
+) -> StandardLoggingGuardrailInformation | None:
+    """The guardrail entry for a trace Bedrock returned because the deployment's ``guardrailConfig`` ran with
+    ``trace: enabled``: every assessment's ``invocationMetrics.usage`` summed and priced like the proxy-side Bedrock
+    guardrail. None when the trace carries no usage, which is what ``trace: disabled`` leaves."""
+    try:
+        guardrail: Final = _JSON_OBJECT.validate_python(trace.get("guardrail"))
+        pairs: Final = _GuardrailTrace.model_validate(guardrail).assessments
+    except ValidationError:
+        return None
+    usage: Final = _summed_usage(tuple(assessment.billable_usage for _, assessment in pairs))
+    if not usage:
+        return None
+    cost_by_unit: Final = bedrock_guardrail_cost_by_unit(usage, aws_region_name)
+    name: Final = guardrail_identifier or pairs[0][0]
+    status: Final[GuardrailStatus] = (
+        "guardrail_intervened" if any(action in _INTERVENING_ACTIONS for action in _actions(guardrail)) else "success"
+    )
+    priced: Final[StandardLoggingGuardrailInformation] = (
+        {"guardrail_cost_by_unit": cost_by_unit} if cost_by_unit is not None else {}
+    )
+    return {
+        "guardrail_name": name,
+        "guardrail_id": name,
+        "guardrail_provider": "bedrock",
+        "guardrail_mode": GuardrailEventHooks.during_call,
+        "guardrail_response": redact_nested_match_and_regex_keys(guardrail),
+        "guardrail_status": status,
+        "guardrail_usage": usage,
+        "guardrail_cost": guardrail_cost_total(cost_by_unit),
+        **priced,
+    }
+
+
+def guardrail_entries(raw: object) -> tuple[object, ...]:
+    """The entries under ``standard_logging_guardrail_information``, which older writers stored as one dict."""
+    single: Final = object_mapping(raw)
+    return (single,) if single is not None else object_items(raw)
+
+
+def with_guardrail_entry(
+    metadata: dict[str, object],  # mutable-ok: the proxy's own request metadata, read later for headers and spend
+    entry: StandardLoggingGuardrailInformation,
+) -> None:
+    """Append ``entry`` to the metadata's ``standard_logging_guardrail_information`` once, so the same trace recorded
+    from the response headers path and again from the success handler bills a single time."""
+    existing: Final = guardrail_entries(metadata.get("standard_logging_guardrail_information"))
+    if entry in existing:
+        return
+    metadata["standard_logging_guardrail_information"] = [*existing, entry]  # rebind-ok: appending here is the job
+
+
 AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"
 
 

litellm/litellm_core_utils/redact_messages.py

@@ -14,7 +14,7 @@ from collections.abc import Mapping, Sequence
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any, Final
 
-from pydantic import TypeAdapter
+from pydantic import JsonValue, TypeAdapter, ValidationError
 
 import litellm
 from litellm.constants import REDACTED_BY_LITELLM
@@ -120,11 +120,55 @@ def _redacted_streaming_response_copy(streaming_response: object) -> object:
     return _redact_streaming_response(redacted_response)
 
 
+BEDROCK_GUARDRAIL_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace")
+_MATCHED_TEXT_KEYS: Final = frozenset({"match", "regex"})
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+
+def _without_matched_text(trace: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    redacted: Final = dict(copy.deepcopy(trace))
+    stack: Final[list[JsonValue]] = [redacted]  # mutable-ok: iterative walk, recursive_detector bans recursion here
+    while stack:
+        node = stack.pop()
+        if isinstance(node, dict):
+            for key in _MATCHED_TEXT_KEYS & node.keys():
+                node[key] = "[REDACTED]"
+            stack.extend(node.values())
+        elif isinstance(node, list):
+            stack.extend(node)
+    return redacted
+
+
+def _redacted_bedrock_trace(trace: object) -> object:
+    try:
+        return _without_matched_text(_JSON_OBJECT.validate_python(trace))
+    except ValidationError:
+        return trace
+
+
+def _redact_bedrock_guardrail_trace(response: object) -> None:
+    """A Bedrock guardrail trace echoes the matched text of its sensitive-information and word findings back."""
+    mapping: Final = _as_mapping(response)
+    if isinstance(response, dict):
+        traces: Final = {
+            field: _redacted_bedrock_trace(value)
+            for field, value in (mapping or {}).items()
+            if field in BEDROCK_GUARDRAIL_TRACE_FIELDS
+        }
+        response.update(traces)
+        return
+    for field, value in ((name, getattr(response, name, None)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS):
+        if value is not None:
+            setattr(response, field, _redacted_bedrock_trace(value))
+
+
 def _redact_streaming_response(streaming_response):
     if hasattr(streaming_response, "choices"):
         for choice in streaming_response.choices:
             _redact_choice_content(choice)
         redact_vertex_ai_metadata_from_logged_object(streaming_response)
+        if isinstance(streaming_response, litellm.ModelResponse):
+            _redact_bedrock_guardrail_trace(streaming_response)
     elif isinstance(streaming_response, litellm.ResponsesAPIResponse):
         return perform_redaction({}, streaming_response)
     elif hasattr(streaming_response, "output"):
@@ -351,6 +395,7 @@ def _redact_standard_logging_object(
             if isinstance(response.get("choices"), list):
                 _redact_model_response_dict_choices(response["choices"], redacted_str)
             redact_vertex_ai_metadata_from_logged_object(response)
+            _redact_bedrock_guardrail_trace(response)
         elif isinstance(response, str):
             standard_logging_object["response"] = redacted_str
         else:
@@ -523,11 +568,13 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
                 for choice in _result.choices:
                     _redact_choice_content(choice)
             redact_vertex_ai_metadata_from_logged_object(_result)
+            _redact_bedrock_guardrail_trace(_result)
         elif isinstance(_result, dict) and "choices" in _result:
             # Handle dict representation of ModelResponse (e.g., from model_dump())
             if _result.get("choices") is not None:
                 _redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM)
             redact_vertex_ai_metadata_from_logged_object(_result)
+            _redact_bedrock_guardrail_trace(_result)
         elif isinstance(_result, dict) and "output" in _result:
             return _redacted_responses_api_response(_result, redact_instructions=redact_streaming_responses)
         elif isinstance(_result, litellm.ResponsesAPIResponse):

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -14,6 +14,7 @@ import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost
 from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_content_from_model_response,  # pyright: ignore[reportUnknownVariableType]  # response param accepts an untyped dict
 )
@@ -1484,6 +1485,35 @@ def _simple_text_part(choices: Sequence[object]) -> str | None:
     return content if isinstance(content, str) else ""
 
 
+_NO_TRACE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+def _carry_bedrock_guardrail_trace(chunks: Sequence[object], response: ModelResponse) -> None:
+    """Bedrock reports a guardrail trace once per stream, on the final metadata chunk, so the built response carries it
+    under the same field as the non-streaming response (``trace`` on Converse, ``amazon-bedrock-trace`` on Invoke)."""
+    traced: Final = map(_bedrock_trace_fields, reversed(chunks))
+    carried: Final = next((fields for fields in traced if fields), _NO_TRACE_FIELDS)
+    for field, trace in carried.items():
+        setattr(response, field, trace)
+
+
+def _record_bedrock_guardrail_trace(
+    chunks: Sequence[object], response: ModelResponse, logging_obj: Optional["Logging"]
+) -> None:
+    _carry_bedrock_guardrail_trace(chunks, response)
+    if logging_obj is not None:
+        logging_obj.record_native_bedrock_guardrail(response)
+
+
+def _bedrock_trace_fields(chunk: object) -> Mapping[str, object]:
+    as_mapping: Final = object_mapping(chunk)
+    raw: Final = (
+        as_mapping.get("provider_specific_fields") if as_mapping else getattr(chunk, "provider_specific_fields", None)
+    )
+    fields: Final[Mapping[str, object]] = object_mapping(raw) or {}
+    return {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields}
+
+
 def stream_chunk_builder(
     chunks: list,
     messages: Sequence | None = None,
@@ -1565,6 +1595,11 @@ def stream_chunk_builder(
                     )
                     break
 
+            _record_bedrock_guardrail_trace(
+                object_items(chunks),  # pyright: ignore[reportUnknownArgumentType]  # chunks is an untyped list param
+                response,
+                logging_obj,
+            )
             _stamp_streaming_usage_cost(usage, response, logging_obj)
             _set_stream_builder_response_cost(
                 response,
@@ -1760,6 +1795,11 @@ def stream_chunk_builder(
             chunks,  # pyright: ignore[reportUnknownArgumentType]  # legacy chunk list is untyped
         )
 
+        _record_bedrock_guardrail_trace(
+            object_items(chunks),  # pyright: ignore[reportUnknownArgumentType]  # chunks is an untyped list param
+            response,
+            logging_obj,
+        )
         processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
         return response
     except Exception as e:

litellm/litellm_core_utils/streaming_handler.py

@@ -1610,6 +1610,8 @@ class CustomStreamWrapper:
         return _ProviderChunkParsed(response_obj)
 
     def chunk_creator(self, chunk: Any):
+        if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields:
+            self.logging_obj.record_native_bedrock_guardrail(chunk.provider_specific_fields)
         if self.custom_llm_provider != "cached_response" and isinstance(
             chunk, (ChatCompletionChunk, ModelResponseStream)
         ):
@@ -1717,6 +1719,8 @@ class CustomStreamWrapper:
                 model_response.usage = Usage(**chunk_usage.model_dump())
 
             ## RETURN ARG
+            if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields:
+                self.copy_model_response_level_provider_specific_fields(chunk, model_response)
             result: Final = self.return_processed_chunk_logic(
                 completion_obj=completion_obj,
                 model_response=model_response,

litellm/llms/anthropic/messages/adapters/handler.py

@@ -806,6 +806,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
                 tool_name_mapping=tool_name_mapping,
                 polyfill_result=polyfill_result,
                 is_async=True,
+                completion_request=completion_kwargs,
                 litellm_logging_obj=logging_obj,
                 safeguards_evaluator=safeguards_evaluator,
             )
@@ -952,6 +953,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
                 tool_name_mapping=tool_name_mapping,
                 polyfill_result=polyfill_result,
                 is_async=False,
+                completion_request=completion_kwargs,
                 litellm_logging_obj=logging_obj,
             )
             if transformed_stream is not None:

litellm/llms/anthropic/messages/adapters/streaming_iterator.py

@@ -22,6 +22,7 @@ from typing import (
 from pydantic import TypeAdapter
 from typing_extensions import assert_never
 
+import litellm
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
 from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError
@@ -41,19 +42,25 @@ from litellm.types.llms.anthropic import (
     UsageDelta,
     UsageIteration,
 )
-from litellm.types.llms.openai import ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock
+from litellm.types.llms.openai import (
+    AllMessageValues,
+    ChatCompletionRedactedThinkingBlock,
+    ChatCompletionThinkingBlock,
+    ChatCompletionToolParam,
+)
 from litellm.types.utils import AdapterCompletionStreamWrapper, ChatCompletionDeltaToolCall, Delta, Function
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
     from litellm.llms.anthropic.messages.safeguards import SafeguardsEvaluator
-    from litellm.types.llms.openai import AllMessageValues
     from litellm.types.utils import ModelResponseStream
 
 
 _STREAMING_DELTA_TYPES: Final = frozenset(get_args(StreamingContentBlockDeltaType))
 _SSE_EVENT: Final = TypeAdapter(Mapping[str, object])
 _THINKING_BLOCK: Final = TypeAdapter(ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock)
+_PROMPT_MESSAGES: Final = TypeAdapter(list[AllMessageValues])
+_PROMPT_TOOLS: Final = TypeAdapter(list[ChatCompletionToolParam])
 
 
 class _UsageDeltaWithIterations(UsageDelta, total=False):
@@ -548,6 +555,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
         iterations_usage: list[UsageIteration] | None = None,
         litellm_logging_obj: "LiteLLMLoggingObject | None" = None,
         safeguards_evaluator: "SafeguardsEvaluator | None" = None,
+        completion_request: Mapping[str, object] | None = None,
     ):
         # Wrap the upstream stream so chunks that carry both content and a
         # finish_reason (fake-streamed providers) are split into two — see
@@ -556,6 +564,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
         super().__init__(combined_chunk_splitter)
         self._combined_chunk_splitter: Final = combined_chunk_splitter
         self.model = model
+        self.completion_request: Final = completion_request
         self._message_id: str = f"msg_{uuid.uuid4()}"
         if litellm_logging_obj is not None:
             litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
@@ -752,25 +761,38 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
 
     def _create_initial_usage_delta(self) -> UsageDelta:
         """
-        Create the initial UsageDelta for the message_start event.
+        Estimate prompt tokens for message_start without consuming upstream chunks.
 
-        Initializes cache token fields (cache_creation_input_tokens, cache_read_input_tokens)
-        to 0 to indicate to clients (like Claude Code) that prompt caching is supported.
-
-        The actual cache token values will be provided in the message_delta event at the
-        end of the stream, since Bedrock Converse API only returns usage data in the final
-        response chunk.
-
-        Returns:
-            UsageDelta with all token counts initialized to 0.
+        Output and cache counts start at zero. Provider usage in message_delta
+        remains authoritative. Failed estimates fall back to zero.
         """
         return UsageDelta(
-            input_tokens=0,
+            input_tokens=self._estimate_initial_input_tokens(),
             output_tokens=0,
             cache_creation_input_tokens=0,
             cache_read_input_tokens=0,
         )
 
+    def _estimate_initial_input_tokens(self) -> int:
+        try:
+            request: Final = self.completion_request or {}
+            messages: Final = _PROMPT_MESSAGES.validate_python(request.get("messages", self.messages))
+            if not messages:
+                return 0
+            materialized_messages: Final = cast(  # cast-ok: JSON-mode serialization preserves validated message types and materializes iterable content
+                list[AllMessageValues], _PROMPT_MESSAGES.dump_python(messages, mode="json")
+            )
+            model: Final = request.get("model", self.model)
+            tools: Final = request.get("tools")
+            return litellm.token_counter(  # pyright: ignore[reportUnknownMemberType]  # token_counter's public signature is partially untyped
+                model=model if isinstance(model, str) else self.model,
+                messages=materialized_messages,
+                tools=_PROMPT_TOOLS.validate_python(tools) if tools is not None else None,
+                use_default_image_token_count=True,
+            )
+        except Exception:  # noqa: BLE001  # prompt estimation must not interrupt the stream
+            return 0
+
     def __next__(self):
         from .transformation import LiteLLMAnthropicMessagesAdapter
 

litellm/llms/anthropic/messages/adapters/transformation.py

@@ -322,6 +322,7 @@ class AnthropicAdapter:
         is_async: bool = True,
         litellm_logging_obj: "LiteLLMLoggingObject | None" = None,
         safeguards_evaluator: "SafeguardsEvaluator | None" = None,
+        completion_request: Mapping[str, object] | None = None,
     ) -> AsyncIterator[bytes] | Iterator[bytes] | None:
         """
         Translate OpenAI streaming response to Anthropic format.
@@ -350,6 +351,7 @@ class AnthropicAdapter:
             iterations_usage=iterations_usage,
             litellm_logging_obj=litellm_logging_obj,
             safeguards_evaluator=safeguards_evaluator,
+            completion_request=completion_request,
         )
         # Return the SSE-wrapped version for proper event formatting.
         if is_async:

litellm/llms/bedrock/chat/converse_transformation.py

@@ -2859,6 +2859,8 @@ class AmazonConverseConfig(BaseConfig):
         # Add "trace" from Bedrock guardrails - if user has opted in to returning it
         if "trace" in completion_response:
             setattr(model_response, "trace", completion_response["trace"])
+            if logging_obj is not None:
+                logging_obj.record_native_bedrock_guardrail(model_response)
 
         # Add service_tier if present in Bedrock response
         # Map Bedrock serviceTier (object) to OpenAI service_tier (string)

litellm/llms/bedrock/chat/invoke_handler.py

@@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache
 from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
 from litellm.litellm_core_utils.core_helpers import map_finish_reason
 from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping
 from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
 from litellm.litellm_core_utils.optional_imports import ensure_optional_import
 from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data
@@ -22,6 +23,7 @@ from litellm.llms.anthropic.chat.handler import (
 from litellm.llms.anthropic.chat.transformation import JsonModeRequest
 from litellm.llms.base_llm import base_model_iterator
 from litellm.llms.base_llm.base_utils import ResponseFormatSchema
+from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -962,8 +964,19 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
             response_format_schema=response_format_schema,
         )
 
-    def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
-        return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
+    def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream:
+        parsed: Final = self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
+        trace: Final = object_mapping(chunk_data.get(BEDROCK_INVOKE_TRACE_FIELD))
+        if trace is None:
+            return parsed
+        return parsed.model_copy(
+            update={
+                "provider_specific_fields": {
+                    **(parsed.provider_specific_fields or {}),
+                    BEDROCK_INVOKE_TRACE_FIELD: trace,
+                }
+            }
+        )
 
     def _stream_end_chunk(self) -> ModelResponseStream | None:
         return self.anthropic_model_response_iterator.cut_stream_chunk()

litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py

@@ -5,6 +5,7 @@ import httpx
 from pydantic import TypeAdapter
 
 from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping
 from litellm.litellm_core_utils.prompt_templates.factory import (
     convert_to_anthropic_image_obj,
 )
@@ -22,6 +23,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
     AmazonInvokeConfig,
 )
 from litellm.llms.bedrock.common_utils import (
+    BEDROCK_INVOKE_TRACE_FIELD,
     apply_bedrock_invoke_structured_output,
     bedrock_chat_uses_native_structured_output,
     bedrock_cleaned_output_schema,
@@ -406,7 +408,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
         api_key: str | None = None,
         json_mode: bool | None = None,
     ) -> ModelResponse:
-        return AnthropicConfig.transform_response(
+        response: Final = AnthropicConfig.transform_response(
             self,
             model=model,
             raw_response=raw_response,
@@ -420,3 +422,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
             api_key=api_key,
             json_mode=json_mode,
         )
+        body: Final = object_mapping(raw_response.json()) or {}
+        trace: Final = object_mapping(body.get(BEDROCK_INVOKE_TRACE_FIELD))
+        if trace is not None:
+            setattr(response, BEDROCK_INVOKE_TRACE_FIELD, trace)
+            logging_obj.record_native_bedrock_guardrail(response)
+        return response

litellm/llms/bedrock/common_utils.py

@@ -45,6 +45,8 @@ if TYPE_CHECKING:
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
     from litellm.types.llms.openai import AllMessageValues
 
+BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace"
+
 
 _ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"
 _OPENAI_FAMILY_MODEL_RE: Final = re.compile(r"(^|[./])openai\.")

litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py

@@ -33,6 +33,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import (
 from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
 from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
     AmazonInvokeConfig,
+    bedrock_invoke_guardrail_headers,
 )
 from litellm.llms.bedrock.common_utils import (
     BedrockError,
@@ -151,8 +152,14 @@ class AmazonAnthropicClaudeMessagesConfig(
         api_key: str | None = None,
         api_base: str | None = None,
     ) -> tuple[dict, str | None]:
+        raw_guardrail_config: Final = litellm_params.get("guardrailConfig")  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]  # litellm_params is an untyped dict param
+        guardrail_headers: Final = (
+            ()
+            if raw_guardrail_config is None
+            else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items())  # pyright: ignore[reportUnknownArgumentType]  # read from the untyped dict param
+        )
         owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
-        return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names), api_base
+        return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names), api_base  # pyright: ignore[reportUnknownArgumentType]  # headers is an untyped dict param
 
     def sign_request(
         self,

litellm/llms/custom_httpx/llm_http_handler.py

@@ -2705,6 +2705,7 @@ class BaseLLMHTTPHandler:
                 raw_response=response,
                 logging_obj=logging_obj,
             )
+            logging_obj.record_native_bedrock_guardrail(initial_response)
 
         return await self._finalize_anthropic_messages_response(
             initial_response=initial_response,

litellm/main.py

@@ -6326,7 +6326,7 @@ def completion(
             response = _complete_vercel_ai_gateway(_dispatch_ctx)
         elif custom_llm_provider == "palm":
             raise ValueError(
-                "Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en"
+                "Palm was decommissioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en"
             )
         elif custom_llm_provider == "vertex_ai_beta" or custom_llm_provider == "gemini":
             response = _complete_vertex_ai_beta(_dispatch_ctx)

litellm/model_prices_and_context_window_backup.json

@@ -28402,7 +28402,6 @@
         "cache_read_input_token_cost_batches": 1.25e-08,
         "cache_read_input_token_cost_flex": 1.25e-08,
         "cache_read_input_token_cost_priority": 4.5e-08,
-        "deprecation_date": "2027-06-28",
         "input_cost_per_image": 0.00028,
         "input_cost_per_token": 2.5e-07,
         "input_cost_per_token_batches": 1.25e-07,
@@ -50878,7 +50877,6 @@
         "cache_read_input_token_cost_batches": 1.25e-08,
         "cache_read_input_token_cost_flex": 1.25e-08,
         "cache_read_input_token_cost_priority": 4.5e-08,
-        "deprecation_date": "2027-06-28",
         "input_cost_per_image": 0.00028,
         "input_cost_per_token": 2.5e-07,
         "input_cost_per_token_batches": 1.25e-07,

litellm/proxy/common_request_processing.py

@@ -6,6 +6,7 @@ import math
 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
 from datetime import datetime
 from functools import lru_cache
+from itertools import chain
 from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
@@ -69,7 +70,11 @@ from litellm.litellm_core_utils.get_supported_openai_params import (
 from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
 from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
+    guardrail_entries,
+    guardrail_information_cost,
+    object_mapping,
+)
 from litellm.litellm_core_utils.llm_response_utils.get_headers import (
     get_response_headers,
 )
@@ -306,6 +311,23 @@ async def close_guarded_stream(stream: object) -> None:
             )
 
 
+def _guardrail_entries_for_headers(
+    request_metadata_bucket: object, logging_obj: LiteLLMLoggingObj
+) -> tuple[object, ...]:
+    """Guardrail entries the proxy hooks wrote on the request plus the ones the call recorded on its own logging
+    object (a Bedrock deployment's native guardrail trace), each once: the hooks copy theirs into the logging object."""
+    litellm_params: Final = object_mapping(logging_obj.model_call_details.get("litellm_params")) or {}
+    buckets: Final = (object_mapping(request_metadata_bucket), object_mapping(litellm_params.get("metadata")))
+    entries: Final = tuple(
+        chain.from_iterable(
+            guardrail_entries(metadata.get("standard_logging_guardrail_information"))
+            for metadata in buckets
+            if metadata is not None
+        )
+    )
+    return tuple(entry for index, entry in enumerate(entries) if entry not in entries[:index])
+
+
 def resolve_litellm_call_id(client_call_id: str | None) -> str:
     if client_call_id is not None and 0 < len(client_call_id) <= MAX_LITELLM_CALL_ID_LENGTH:
         return client_call_id
@@ -3238,9 +3260,9 @@ class ProxyBaseLLMRequestProcessing:
             if is_unbilled_non_inference_call_from_params(logging_obj.call_type, logging_obj.litellm_params, response)
             else computed_cost_for_headers
         )
-        _, request_metadata_bucket = get_or_create_metadata_bucket(self.data)
+        request_metadata_bucket: Final[object] = get_or_create_metadata_bucket(self.data)[1]  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]  # self.data is an untyped dict
         guardrail_cost_for_headers: Final = guardrail_information_cost(
-            request_metadata_bucket.get("standard_logging_guardrail_information")
+            _guardrail_entries_for_headers(request_metadata_bucket, logging_obj)  # pyright: ignore[reportUnknownArgumentType]  # bucket comes from the untyped request dict
         )
         response_cost_for_headers: Final = (
             (llm_cost_for_headers if isinstance(llm_cost_for_headers, (int, float)) else 0.0)

litellm/proxy/db/prisma_client.py

@@ -1126,6 +1126,18 @@ class PrismaManager:
             return False
         return ProxyExtrasDBManager.build_request_log_indexes()
 
+    @staticmethod
+    def migration_data_check_report(use_migrate: bool = True) -> str | None:
+        """The read-only pre-migration row check on its own, for `--check_migrations_against_existing_data`: the
+        refusal the migration step would print, or None when the pending schema changes
+        can be applied to the rows the database has."""
+        from litellm_proxy_extras.utils import (
+            ProxyExtrasDBManager,  # pyright: ignore[reportUnknownVariableType]  # litellm-proxy-extras is an optional install
+        )
+
+        refusal: Final[object] = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # litellm-proxy-extras is an optional install
+        return refusal if isinstance(refusal, str) else None
+
     @staticmethod
     def start_request_log_index_build() -> None:
         """Build the request-log indexes on a daemon thread, for a serving proxy that ran the

litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py

@@ -28,6 +28,7 @@ from litellm.llms.anthropic.messages.stream_assembly import (
     collapse_pure_text_chunks,
     split_sse_chunk_into_events,
 )
+from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD
 from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
 from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
@@ -70,6 +71,17 @@ def _event_payload(event: str) -> Mapping[str, object] | None:
         return None
 
 
+def _bedrock_trace_payload(all_chunks: Sequence[str | bytes]) -> Mapping[str, object] | None:
+    payloads: Final = (
+        _event_payload(event)
+        for chunk in reversed(all_chunks)
+        for event in reversed(split_sse_chunk_into_events(chunk))
+    )  # comprehension-ok: chunk then event, read newest first
+    return next(
+        (payload for payload in payloads if payload is not None and BEDROCK_INVOKE_TRACE_FIELD in payload), None
+    )
+
+
 def _message_delta_usages(all_chunks: Sequence[str | bytes]) -> Iterator[object]:
     for chunk in reversed(all_chunks):
         yield from (
@@ -141,6 +153,7 @@ class AnthropicPassthroughLoggingHandler:
             json_mode=False,
             litellm_params={},
         )
+        logging_obj.record_native_bedrock_guardrail(response_body)  # pyright: ignore[reportUnknownArgumentType]  # response_body is an untyped dict param
 
         kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
             litellm_model_response=litellm_model_response,
@@ -546,6 +559,7 @@ class AnthropicPassthroughLoggingHandler:
                 "result": None,
                 "kwargs": {},
             }
+        litellm_logging_obj.record_native_bedrock_guardrail(_bedrock_trace_payload(all_chunks))
         kwargs: Final = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
             litellm_model_response=complete_streaming_response,
             model=model,

litellm/proxy/proxy_cli.py

@@ -218,6 +218,23 @@ class ProxyInitializationHelpers:
 
     _run_health_check = run_health_check
 
+    @staticmethod
+    def run_migration_data_check(use_migrate: bool) -> int:
+        """`--check_migrations_against_existing_data`: the pre-migration row check against DATABASE_URL, nothing applied.
+        0 when the pending schema changes fit the rows, 2 with the refusal when they do not."""
+        if os.getenv("DATABASE_URL") is None:
+            raise click.UsageError("--check_migrations_against_existing_data requires DATABASE_URL")
+        from litellm.proxy.db.prisma_client import PrismaManager
+
+        refusal: Final[str | None] = PrismaManager.migration_data_check_report(use_migrate=use_migrate)  # pyright: ignore[reportUnknownVariableType]  # litellm-proxy-extras is an optional install
+        if refusal is None:
+            print("LiteLLM Proxy: the pending migrations and indexes can be applied to the rows this database has.")
+            return 0
+        print(
+            f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {refusal}\033[0m", file=sys.stderr, flush=True
+        )
+        return 2
+
     @staticmethod
     def run_config_validation(config: str | None) -> None:
         if config is None:
@@ -978,6 +995,13 @@ class ProxyInitializationHelpers:
     default=False,
     help="Load and validate the config file (including mcp_servers) without starting the server, then exit. Exit code 1 on any config error.",
 )
+@click.option(
+    "--check_migrations_against_existing_data",
+    is_flag=True,
+    default=False,
+    help="Check, read-only, that the rows in DATABASE_URL can take this version's pending migrations and indexes, "
+    "print what would block them, then exit without applying anything. Exit code 2 when something blocks.",
+)
 @click.option(
     "--keepalive_timeout",
     default=None,
@@ -1127,6 +1151,7 @@ def run_server(
     use_prisma_db_push: bool,
     skip_server_startup,
     validate_config: bool,
+    check_migrations_against_existing_data: bool,
     keepalive_timeout,
     timeout_worker_healthcheck,
     max_requests_before_restart: int | None,
@@ -1172,6 +1197,8 @@ def run_server(
     if validate_config is True:
         ProxyInitializationHelpers.run_config_validation(config)
         return
+    if check_migrations_against_existing_data is True:
+        sys.exit(ProxyInitializationHelpers.run_migration_data_check(use_migrate=not use_prisma_db_push))
     if enforce_prisma_migration_check:
         print(
             "\033[1;33mLiteLLM Proxy: --enforce_prisma_migration_check is "

litellm/responses/litellm_completion_transformation/handler.py

@@ -133,7 +133,6 @@ class LiteLLMCompletionTransformationHandler:
         litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion(
             **acompletion_args,
         )
-
         if isinstance(litellm_completion_response, ModelResponse):
             responses_api_response: Final[ResponsesAPIResponse] = (
                 LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(

litellm/responses/litellm_completion_transformation/transformation.py

@@ -2062,7 +2062,19 @@ class LiteLLMCompletionResponsesConfig:
         - ItemReference
         """
         if input_item.get("type") == "web_search_call":
-            search: Final = ResponseFunctionWebSearch.model_validate(input_item)
+            try:
+                search: Final = ResponseFunctionWebSearch.model_validate(input_item)
+            except ValidationError as exc:
+                errors: Final = exc.errors(include_url=False, include_input=False)
+                if len(errors) != 1 or errors[0]["type"] != "missing" or errors[0]["loc"] != ("action",):
+                    raise
+                replay: Final = ResponseFunctionWebSearch.model_construct(None, **input_item)
+                return [
+                    GenericChatCompletionMessage(
+                        role="assistant",
+                        content="Hosted web search: " + replay.model_dump_json(exclude_none=True),
+                    )
+                ]
             return [
                 GenericChatCompletionMessage(
                     role="assistant",

model_prices_and_context_window.json

@@ -28402,7 +28402,6 @@
         "cache_read_input_token_cost_batches": 1.25e-08,
         "cache_read_input_token_cost_flex": 1.25e-08,
         "cache_read_input_token_cost_priority": 4.5e-08,
-        "deprecation_date": "2027-06-28",
         "input_cost_per_image": 0.00028,
         "input_cost_per_token": 2.5e-07,
         "input_cost_per_token_batches": 1.25e-07,
@@ -50878,7 +50877,6 @@
         "cache_read_input_token_cost_batches": 1.25e-08,
         "cache_read_input_token_cost_flex": 1.25e-08,
         "cache_read_input_token_cost_priority": 4.5e-08,
-        "deprecation_date": "2027-06-28",
         "input_cost_per_image": 0.00028,
         "input_cost_per_token": 2.5e-07,
         "input_cost_per_token_batches": 1.25e-07,

scripts/generate_provider_endpoints_support.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections.abc import Iterator, Mapping, Sequence
+from itertools import chain
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+from pydantic import JsonValue, TypeAdapter
+
+ROOT: Final = Path(__file__).resolve().parents[1]
+COST_MAP_PATH: Final = ROOT / "model_prices_and_context_window.json"
+SOURCE_PATH: Final = ROOT / "provider_endpoints_support.json"
+OUTPUT_PATHS: Final = (
+    SOURCE_PATH,
+    ROOT / "litellm/provider_endpoints_support_backup.json",
+)
+JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+MODE_TO_FLAG: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "chat": "chat_completions",
+        "responses": "responses",
+        "completion": "text_completion",
+        "embedding": "embeddings",
+        "image_generation": "image_generations",
+        "image_edit": "image_edits",
+        "audio_transcription": "audio_transcriptions",
+        "audio_speech": "audio_speech",
+        "rerank": "rerank",
+        "ocr": "ocr",
+        "search": "search",
+        "moderation": "moderations",
+        "realtime": "realtime",
+        "video_generation": "video_generations",
+    }
+)
+SUPPORTED_ENDPOINT_TO_FLAG: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "/v1/chat/completions": "chat_completions",
+        "/v1/responses": "responses",
+        "/v1/completions": "text_completion",
+        "/v1/embeddings": "embeddings",
+        "/v1/images/generations": "image_generations",
+        "/v1/images/edits": "image_edits",
+        "/v1/audio/transcriptions": "audio_transcriptions",
+        "/v1/audio/speech": "audio_speech",
+        "/v1/ocr": "ocr",
+        "/v1/realtime": "realtime",
+        "/v1/videos": "video_generations",
+        "/v1/batch": "batches",
+        "/v1/messages": "messages",
+        "/v1beta/interactions": "interactions",
+        "/v1/systemone": "systemone",
+    }
+)
+
+
+def _json_object(value: JsonValue, name: str) -> Mapping[str, JsonValue]:
+    if not isinstance(value, dict):
+        raise ValueError(f"{name} must be a JSON object")
+    return value
+
+
+def _provider_slug(provider: str) -> str:
+    if provider.startswith("vertex_ai-"):
+        return "vertex_ai"
+    if provider == "fireworks_ai-embedding-models":
+        return "fireworks_ai"
+    if provider == "bedrock_converse":
+        return "bedrock"
+    if provider == "text-completion-inception":
+        return "inception"
+    return provider
+
+
+def _model_provider_slug(model: JsonValue) -> str | None:
+    if not isinstance(model, dict):
+        return None
+    provider: Final = model.get("litellm_provider")
+    if not isinstance(provider, str) or not provider:
+        return None
+    return _provider_slug(provider)
+
+
+def _model_flags(model: Mapping[str, JsonValue]) -> tuple[str, ...]:
+    mode: Final = model.get("mode")
+    mode_flag: Final = MODE_TO_FLAG.get(mode) if isinstance(mode, str) else None
+    supported_endpoints: Final = model.get("supported_endpoints")
+    endpoint_flags: Final = (
+        tuple(
+            SUPPORTED_ENDPOINT_TO_FLAG[endpoint]
+            for endpoint in supported_endpoints
+            if isinstance(endpoint, str) and endpoint in SUPPORTED_ENDPOINT_TO_FLAG
+        )
+        if isinstance(supported_endpoints, list)
+        else ()
+    )
+    return ((mode_flag,) if mode_flag is not None else ()) + endpoint_flags
+
+
+def _bridge_field(endpoint: JsonValue) -> str | None:
+    if not isinstance(endpoint, dict) or endpoint.get("bridges_to_chat_completion") is not True:
+        return None
+    provider_json_field: Final = endpoint.get("provider_json_field")
+    return provider_json_field if isinstance(provider_json_field, str) else None
+
+
+def _bridged_fields(endpoints: Mapping[str, JsonValue]) -> tuple[str, ...]:
+    return tuple(
+        dict.fromkeys(
+            provider_json_field
+            for endpoint in endpoints.values()
+            if (provider_json_field := _bridge_field(endpoint)) is not None
+        )
+    )
+
+
+def _derived_keys(bridged_fields: tuple[str, ...]) -> tuple[str, ...]:
+    return tuple(dict.fromkeys((*MODE_TO_FLAG.values(), *SUPPORTED_ENDPOINT_TO_FLAG.values(), *bridged_fields)))
+
+
+def _model_evidence(
+    model: JsonValue, bridged_fields: tuple[str, ...]
+) -> Iterator[tuple[str, str]]:
+    provider: Final = _model_provider_slug(model)
+    if provider is None or not isinstance(model, dict):
+        return
+    flags: Final = _model_flags(model)
+    yield from ((provider, flag) for flag in flags)
+    if "chat_completions" not in flags:
+        return
+    yield from ((provider, flag) for flag in bridged_fields)
+
+
+def _provider_evidence(
+    cost_map: Mapping[str, JsonValue], bridged_fields: tuple[str, ...]
+) -> Iterator[tuple[str, str]]:
+    return chain.from_iterable(
+        _model_evidence(model, bridged_fields)
+        for model_name, model in cost_map.items()
+        if model_name != "sample_spec"
+    )
+
+
+def _provider_slugs(cost_map: Mapping[str, JsonValue]) -> frozenset[str]:
+    return frozenset(
+        provider_slug
+        for model_name, model in cost_map.items()
+        if model_name != "sample_spec" and (provider_slug := _model_provider_slug(model)) is not None
+    )
+
+
+def _generated_provider(
+    slug: str,
+    value: JsonValue,
+    derived_keys: tuple[str, ...],
+    derived_key_set: frozenset[str],
+    evidence: frozenset[tuple[str, str]],
+) -> dict[str, JsonValue]:
+    provider: Final = _json_object(value, f"provider {slug}")
+    current_flags_value: Final = provider.get("endpoints")
+    current_flags: Final = (
+        _json_object(current_flags_value, f"provider {slug} endpoints")
+        if current_flags_value is not None
+        else {}
+    )
+    derived_flags: Final = {flag: (slug, flag) in evidence for flag in derived_keys}
+    handwritten_flags: Final = {
+        flag: value for flag, value in current_flags.items() if flag not in derived_key_set
+    }
+    return {
+        **provider,
+        "endpoints": {**derived_flags, **handwritten_flags},
+    }
+
+
+def generate_provider_endpoints_support(
+    cost_map: Mapping[str, JsonValue],
+    current_json: Mapping[str, JsonValue],
+) -> dict[str, JsonValue]:
+    current_providers: Final = _json_object(
+        current_json.get("providers"), "provider_endpoints_support providers"
+    )
+    endpoints: Final = _json_object(current_json.get("endpoints"), "provider_endpoints_support endpoints")
+    bridged_fields: Final = _bridged_fields(endpoints)
+    derived_keys: Final = _derived_keys(bridged_fields)
+    derived_key_set: Final = frozenset(derived_keys)
+    provider_slugs: Final = _provider_slugs(cost_map)
+    missing_slugs: Final = tuple(sorted(provider_slugs - frozenset(current_providers)))
+    if missing_slugs:
+        raise ValueError(f"Cost map providers missing from provider_endpoints_support.json: {', '.join(missing_slugs)}")
+    evidence: Final = frozenset(_provider_evidence(cost_map, bridged_fields))
+    providers: Final = {
+        slug: _generated_provider(slug, value, derived_keys, derived_key_set, evidence)
+        for slug, value in current_providers.items()
+    }
+    return {**current_json, "providers": providers}
+
+
+def generate_files(
+    cost_map: Mapping[str, JsonValue],
+    current_json: Mapping[str, JsonValue],
+    output_paths: tuple[Path, ...],
+    check: bool,
+) -> int:
+    generated: Final = generate_provider_endpoints_support(cost_map, current_json)
+    content: Final = json.dumps(generated, indent=2, ensure_ascii=False) + "\n"
+    stale_paths: Final = tuple(
+        path for path in output_paths if not path.exists() or path.read_text(encoding="utf-8") != content
+    )
+    if check:
+        if not stale_paths:
+            return 0
+        sys.stderr.write(
+            "provider endpoint support is stale in "
+            + ", ".join(str(path) for path in stale_paths)
+            + "\nRun `uv run --no-sync python scripts/generate_provider_endpoints_support.py` to regenerate it\n"
+        )
+        return 1
+    for path in output_paths:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(content, encoding="utf-8")
+    return 0
+
+
+def main(argv: Sequence[str] | None = None, *, repo_root: Path = ROOT) -> int:
+    parser: Final = argparse.ArgumentParser(description="Generate provider endpoint support from the model cost map")
+    parser.add_argument("--check", action="store_true", help="fail if either generated file is stale")
+    args: Final = parser.parse_args(argv)
+    cost_map: Final = JSON_OBJECT.validate_json((repo_root / COST_MAP_PATH.name).read_bytes())
+    current_json: Final = JSON_OBJECT.validate_json((repo_root / SOURCE_PATH.name).read_bytes())
+    return generate_files(
+        cost_map,
+        current_json,
+        (
+            repo_root / SOURCE_PATH.name,
+            repo_root / OUTPUT_PATHS[1].relative_to(ROOT),
+        ),
+        args.check,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/code_coverage_tests/test_generate_provider_endpoints_support.py

@@ -0,0 +1,179 @@
+import json
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Final
+
+import pytest
+from pydantic import JsonValue
+
+from scripts.generate_provider_endpoints_support import (
+    generate_provider_endpoints_support,
+    main,
+)
+
+ENDPOINTS: Final = {
+    "a2a": {"provider_json_field": "a2a", "bridges_to_chat_completion": True},
+    "messages": {"provider_json_field": "messages", "bridges_to_chat_completion": True},
+    "responses": {"provider_json_field": "responses", "bridges_to_chat_completion": True},
+    "generateContent": {"provider_json_field": "generateContent", "bridges_to_chat_completion": True},
+    "text_completion": {"provider_json_field": "text_completion", "bridges_to_chat_completion": True},
+    "count_tokens": {"provider_json_field": "count_tokens"},
+}
+
+
+def _current_json(providers: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+    return {
+        "_comment": "retained",
+        "_schema": {"version": 1},
+        "providers": dict(providers),
+        "endpoints": ENDPOINTS,
+    }
+
+
+def _provider(flags: Mapping[str, JsonValue] | None = None) -> dict[str, JsonValue]:
+    return {
+        "display_name": "Example (`example`)",
+        "url": "https://docs.litellm.ai/docs/providers/example",
+        "endpoints": dict(flags or {}),
+    }
+
+
+def _flags(generated: Mapping[str, JsonValue], slug: str) -> Mapping[str, JsonValue]:
+    providers: Final = generated["providers"]
+    assert isinstance(providers, dict)
+    provider: Final = providers[slug]
+    assert isinstance(provider, dict)
+    flags: Final = provider["endpoints"]
+    assert isinstance(flags, dict)
+    return flags
+
+
+def test_mode_derives_chat_flag() -> None:
+    cost_map: Final = {"example/model": {"litellm_provider": "example", "mode": "chat"}}
+    current_json: Final = _current_json({"example": _provider()})
+
+    generated: Final = generate_provider_endpoints_support(cost_map, current_json)
+
+    assert _flags(generated, "example")["chat_completions"] is True
+
+
+def test_supported_endpoint_derives_embedding_flag() -> None:
+    cost_map: Final = {
+        "example/model": {
+            "litellm_provider": "example",
+            "mode": "completion",
+            "supported_endpoints": ["/v1/embeddings"],
+        }
+    }
+    current_json: Final = _current_json({"example": _provider()})
+
+    generated: Final = generate_provider_endpoints_support(cost_map, current_json)
+
+    assert _flags(generated, "example")["embeddings"] is True
+
+
+def test_chat_models_enable_bridge_flags_but_other_models_do_not() -> None:
+    cost_map: Final = {
+        "example/chat": {"litellm_provider": "has_chat", "mode": "chat"},
+        "example/endpoint-chat": {
+            "litellm_provider": "endpoint_chat",
+            "mode": "image_generation",
+            "supported_endpoints": ["/v1/chat/completions"],
+        },
+        "example/image": {"litellm_provider": "no_chat", "mode": "image_generation"},
+    }
+    current_json: Final = _current_json(
+        {
+            "has_chat": _provider(),
+            "endpoint_chat": _provider(),
+            "no_chat": _provider(),
+        }
+    )
+
+    generated: Final = generate_provider_endpoints_support(cost_map, current_json)
+
+    assert _flags(generated, "has_chat")["a2a"] is True
+    assert _flags(generated, "has_chat")["generateContent"] is True
+    assert _flags(generated, "endpoint_chat")["a2a"] is True
+    assert _flags(generated, "no_chat")["a2a"] is False
+    assert _flags(generated, "no_chat")["generateContent"] is False
+
+
+def test_unbacked_derived_flags_reset_and_handwritten_flags_remain() -> None:
+    cost_map: Final = {"example/model": {"litellm_provider": "example", "mode": "image_generation"}}
+    current_json: Final = _current_json(
+        {
+            "example": _provider({"responses": True, "count_tokens": True, "files": False}),
+            "without_models": _provider({"responses": True}),
+        }
+    )
+
+    generated: Final = generate_provider_endpoints_support(cost_map, current_json)
+    flags: Final = _flags(generated, "example")
+
+    assert flags["responses"] is False
+    assert flags["count_tokens"] is True
+    assert flags["files"] is False
+    assert _flags(generated, "without_models")["chat_completions"] is False
+
+
+def test_provider_aliases_share_their_mapped_slugs() -> None:
+    cost_map: Final = {
+        "vertex/model-a": {"litellm_provider": "vertex_ai-chat", "mode": "chat"},
+        "vertex/model-b": {"litellm_provider": "vertex_ai-embedding", "mode": "embedding"},
+        "bedrock/model": {"litellm_provider": "bedrock_converse", "mode": "responses"},
+    }
+    current_json: Final = _current_json({"vertex_ai": _provider(), "bedrock": _provider()})
+
+    generated: Final = generate_provider_endpoints_support(cost_map, current_json)
+
+    assert _flags(generated, "vertex_ai")["chat_completions"] is True
+    assert _flags(generated, "vertex_ai")["embeddings"] is True
+    assert _flags(generated, "bedrock")["responses"] is True
+
+
+def test_cost_map_provider_without_file_entry_is_an_error() -> None:
+    cost_map: Final = {"example/model": {"litellm_provider": "missing_provider", "mode": "chat"}}
+    current_json: Final = _current_json({"example": _provider()})
+
+    with pytest.raises(ValueError, match="missing_provider"):
+        generate_provider_endpoints_support(cost_map, current_json)
+
+
+def test_check_reports_stale_output_without_writing(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None:
+    cost_map: Final = {"example/model": {"litellm_provider": "example", "mode": "chat"}}
+    current_json: Final = _current_json({"example": _provider()})
+    (tmp_path / "model_prices_and_context_window.json").write_text(json.dumps(cost_map), encoding="utf-8")
+    source_path: Final = tmp_path / "provider_endpoints_support.json"
+    source_content: Final = json.dumps(current_json)
+    source_path.write_text(source_content, encoding="utf-8")
+
+    result: Final = main(["--check"], repo_root=tmp_path)
+
+    assert result == 1
+    assert source_path.read_text(encoding="utf-8") == source_content
+    assert not (tmp_path / "litellm/provider_endpoints_support_backup.json").exists()
+    error_output: Final = capsys.readouterr().err
+    assert "provider endpoint support is stale" in error_output
+    assert "scripts/generate_provider_endpoints_support.py" in error_output
+
+
+def test_check_reports_stale_backup_without_writing(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None:
+    cost_map: Final = {"example/model": {"litellm_provider": "example", "mode": "chat"}}
+    current_json: Final = _current_json({"example": _provider()})
+    generated: Final = generate_provider_endpoints_support(cost_map, current_json)
+    generated_content: Final = json.dumps(generated, indent=2, ensure_ascii=False) + "\n"
+    source_path: Final = tmp_path / "provider_endpoints_support.json"
+    backup_path: Final = tmp_path / "litellm/provider_endpoints_support_backup.json"
+    backup_content: Final = "stale\n"
+    (tmp_path / "model_prices_and_context_window.json").write_text(json.dumps(cost_map), encoding="utf-8")
+    source_path.write_text(generated_content, encoding="utf-8")
+    backup_path.parent.mkdir()
+    backup_path.write_text(backup_content, encoding="utf-8")
+
+    result: Final = main(["--check"], repo_root=tmp_path)
+
+    assert result == 1
+    assert source_path.read_text(encoding="utf-8") == generated_content
+    assert backup_path.read_text(encoding="utf-8") == backup_content
+    assert str(backup_path) in capsys.readouterr().err

tests/integration/database/migrations/test_request_log_indexes.py

@@ -1,5 +1,6 @@
 import os
 import queue
+import secrets
 import shutil
 import subprocess
 import sys
@@ -14,7 +15,7 @@ from typing import Final
 
 import psycopg
 import pytest
-from litellm_proxy_extras import request_log_indexes
+from litellm_proxy_extras import migration_data_checks, request_log_indexes
 from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, migration_lock
 from litellm_proxy_extras.migration_recovery import roll_back_failed_inert_migration
 from litellm_proxy_extras.request_log_indexes import (
@@ -408,6 +409,12 @@ def test_a_migration_without_a_failed_row_is_not_touched(partitioned_database: s
     assert _ledger(partitioned_database) == before
 
 
+def _backend_pid(connection: "psycopg.Connection[tuple[object, ...]]") -> int:
+    row: Final = connection.execute("SELECT pg_backend_pid()").fetchone()
+    assert row is not None and isinstance(row[0], int), row
+    return row[0]
+
+
 def _pin_a_snapshot_on(database_url: str, table: str) -> "psycopg.Connection[tuple[object, ...]]":
     pin: Final = psycopg.connect(database_url)
     pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
@@ -532,7 +539,9 @@ def test_an_invalid_child_index_left_by_an_interrupted_build_is_rebuilt_and_atta
 
 
 @requires_db
-def test_partitioning_a_table_the_index_job_already_covered_moves_its_call_id_index_aside(scratch_database: str) -> None:
+def test_partitioning_a_table_the_index_job_already_covered_moves_its_call_id_index_aside(
+    scratch_database: str,
+) -> None:
     """The runbook renames the Prisma indexes off the legacy table so the partitioned table can take their
     names. The call-id index is built by the index job, not Prisma, and a deployment that partitions after
     that job ran kept the name on the legacy table, so every later job run reported it as on another table."""
@@ -541,7 +550,9 @@ def test_partitioning_a_table_the_index_job_already_covered_moves_its_call_id_in
 
     _partition_spend_logs(scratch_database)
 
-    assert _indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_litellm_call_id_idx") == "LiteLLM_SpendLogs_legacy"
+    assert (
+        _indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_litellm_call_id_idx") == "LiteLLM_SpendLogs_legacy"
+    )
     assert ensure_request_log_indexes(scratch_database, "public") is True
     _assert_index_covers_every_partition(scratch_database, CALL_ID_INDEX, "litellm_call_id_idx")
 
@@ -737,7 +748,7 @@ def test_inserts_are_never_held_back_while_attach_partition_waits_for_a_reader_o
         reader.execute(
             sql.SQL('SELECT count(*) FROM {} WHERE "litellm_call_id" IS NULL').format(sql.Identifier(partition))
         ).fetchone()
-        reader_pid: Final = reader.execute("SELECT pg_backend_pid()").fetchone()[0]
+        reader_pid: Final = _backend_pid(reader)
         with psycopg.connect(partitioned_database, autocommit=True) as inspector:
             child_lock: Final = inspector.execute(
                 "SELECT 1 FROM pg_locks WHERE pid = %s AND relation = to_regclass(%s) "
@@ -800,7 +811,7 @@ def test_an_attach_that_never_gets_its_lock_is_left_for_the_next_index_build(
         reader.execute(
             sql.SQL('SELECT count(*) FROM {} WHERE "litellm_call_id" IS NULL').format(sql.Identifier(partition))
         ).fetchone()
-        reader_pid: Final = reader.execute("SELECT pg_backend_pid()").fetchone()[0]
+        reader_pid: Final = _backend_pid(reader)
         with psycopg.connect(partitioned_database, autocommit=True) as inspector:
             child_lock: Final = inspector.execute(
                 "SELECT 1 FROM pg_locks WHERE pid = %s AND relation = to_regclass(%s) "
@@ -929,3 +940,183 @@ def test_a_second_copy_of_a_managed_index_is_reported_with_its_drop_statement_an
 
     assert 'remove it with: DROP INDEX CONCURRENTLY "public"."call_id_copy"' in caplog.text
     assert _index_oids(scratch_database)["call_id_copy"] == copy
+
+
+def _insert_spend_log_with_key(conn: "psycopg.Connection[tuple[object, ...]]", request_id: str, api_key: str) -> None:
+    conn.execute(
+        'INSERT INTO "LiteLLM_SpendLogs" ("request_id", "call_type", "startTime", "endTime", "api_key") '
+        "VALUES (%s, %s, %s, %s, %s)",
+        (request_id, "acompletion", "2025-06-03 14:02:11", "2025-06-03 14:02:13", api_key),
+    )
+
+
+def _legacy_bearer_token() -> str:
+    """A raw JWT-sized bearer token the way a pre-v1.99.0 proxy stored it: random, so
+    Postgres cannot compress it into the 2704-byte btree entry limit."""
+    return "Bearer " + secrets.token_hex(1600)
+
+
+def _hash_sql_from(refusal: str) -> str:
+    return refusal[refusal.index('UPDATE "LiteLLM_SpendLogs"') : refusal.index("{64}$';") + len("{64}$';")]
+
+
+@requires_db
+@RESOLVERS
+@pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",))
+def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_keys_are_too_long_for_the_index(
+    scratch_database: str, use_v2_resolver: bool
+) -> None:
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        _insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token())
+        _insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token())
+        _insert_spend_log_with_key(conn, "hashed-1", "a" * 64)
+    ledger_before: Final = _ledger(scratch_database)
+
+    with pytest.raises(RuntimeError) as refused:
+        ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver)
+
+    refusal: Final = str(refused.value)
+    assert (
+        '"LiteLLM_SpendLogs".api_key, startTime: 2 row(s) with an index entry over 2704 bytes, largest 3232 bytes '
+        "(api_key up to 3207 bytes)"
+    ) in refusal, refusal
+    assert f'blocks: CREATE INDEX "{API_KEY_INDEX}"' in refusal, refusal
+    assert _ledger(scratch_database) == ledger_before, "a refused job must not apply a single migration"
+    assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx")
+
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute(_hash_sql_from(refusal))  # pyright: ignore[reportArgumentType]  # the refusal's own SQL
+    assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
+    assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
+    assert len(_ledger(scratch_database)) > len(ledger_before)
+
+
+@requires_db
+@pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",))
+def test_a_serving_proxy_index_build_that_meets_too_long_keys_reports_the_rows_instead_of_the_postgres_error(
+    scratch_database: str, caplog: pytest.LogCaptureFixture
+) -> None:
+    assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) is True
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        _insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token())
+
+    with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
+        built: Final = ensure_request_log_indexes(scratch_database, "public")
+
+    assert built is False
+    assert f"Not building index {API_KEY_INDEX}" in caplog.text, caplog.text
+    assert '"LiteLLM_SpendLogs".api_key, startTime: 1 row(s) with an index entry over 2704 bytes' in caplog.text, (
+        caplog.text
+    )
+    assert "exceeds btree version 4 maximum" not in caplog.text, caplog.text
+    assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx")
+
+
+_PROBE_STATEMENTS: Final = (
+    'CREATE UNIQUE INDEX "probe_b_key" ON "preflight_probe" ("b")',
+    'ALTER TABLE "preflight_probe" ALTER COLUMN "b" SET NOT NULL',
+    'ALTER TABLE "preflight_probe" ADD COLUMN "c" INTEGER NOT NULL',
+    'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_r_fkey" FOREIGN KEY ("r") REFERENCES "preflight_parent"("id")',
+    'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_a_check" CHECK ("a" <> \'bad\')',
+    'CREATE INDEX "probe_a_idx" ON "preflight_probe" ("a")',
+)
+
+
+@requires_db
+def test_each_check_kind_counts_exactly_the_rows_postgres_would_reject(scratch_database: str) -> None:
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute('CREATE TABLE "preflight_parent" ("id" INTEGER PRIMARY KEY)')
+        conn.execute('CREATE TABLE "preflight_probe" ("a" TEXT, "b" INTEGER, "r" INTEGER)')
+        conn.execute('INSERT INTO "preflight_parent" VALUES (1)')
+        conn.execute(
+            'INSERT INTO "preflight_probe" VALUES (%s, 1, 1), (%s, 1, 1), (%s, NULL, 2), (%s, 2, NULL), (%s, 3, 1)',
+            ("ok", "ok", "bad", "bad", secrets.token_hex(1400)),
+        )
+        checks: Final = tuple(
+            check
+            for statement in _PROBE_STATEMENTS
+            for check in migration_data_checks.checks_for_statement(statement, "probe")
+        )
+        findings: Final = migration_data_checks.find_blockers(conn, "public", checks)
+        for statement in _PROBE_STATEMENTS:
+            with conn.transaction(force_rollback=True), pytest.raises(psycopg.Error):
+                conn.execute(statement)  # pyright: ignore[reportArgumentType]  # test SQL
+
+    counted: Final = {
+        (f.check.kind, f.check.columns): f.rows for f in findings if isinstance(f, migration_data_checks.BlockedChange)
+    }
+    assert counted == {
+        ("duplicates", ("b",)): 1,
+        ("nulls", ("b",)): 1,
+        ("any_row", ("c",)): 5,
+        ("orphans", ("r",)): 1,
+        ("check_violations", ()): 2,
+        ("too_long", ("a",)): 1,
+    }
+
+
+@requires_db
+@pytest.mark.parametrize(
+    ("index_columns", "s_value", "t_value", "longest_fitting_key", "described_as"),
+    (
+        (("k", "t"), None, "2025-06-03 14:02:11", 2684, "k up to 2685 bytes"),
+        (("s", "k"), "q" * 20, None, 2668, "s up to 20 bytes, k up to 2669 bytes"),
+        (("s", "k"), None, None, 2684, "k up to 2685 bytes"),
+    ),
+    ids=("key then timestamp", "20-byte text then key", "NULL then key"),
+)
+def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does(
+    scratch_database: str,
+    index_columns: tuple[str, ...],
+    s_value: str | None,
+    t_value: str | None,
+    longest_fitting_key: int,
+    described_as: str,
+) -> None:
+    """One row's key is the longest that still fits the index entry, the other's is one byte
+    longer. Where that line falls depends on the other indexed column, its alignment and the
+    null bitmap, so the same two rows go through the check and through Postgres."""
+    create_index: Final = sql.SQL('CREATE INDEX "entry_probe_idx" ON "entry_probe" ({})').format(
+        sql.SQL(", ").join(sql.Identifier(column) for column in index_columns)
+    )
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute('CREATE TABLE "entry_probe" ("s" TEXT, "k" TEXT, "t" TIMESTAMP(3))')
+        conn.execute(
+            'INSERT INTO "entry_probe" ("s", "k", "t") VALUES (%s, %s, %s), (%s, %s, %s)',
+            (
+                s_value,
+                secrets.token_hex(longest_fitting_key)[:longest_fitting_key],
+                t_value,
+                s_value,
+                secrets.token_hex(longest_fitting_key)[: longest_fitting_key + 1],
+                t_value,
+            ),
+        )
+        findings: Final = migration_data_checks.find_blockers(
+            conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
+        )
+        with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="exceeds btree version 4 maximum 2704"):
+            conn.execute(create_index)
+        conn.execute('DELETE FROM "entry_probe" WHERE octet_length("k") > %s', (longest_fitting_key,))
+        conn.execute(create_index)
+
+    (finding,) = findings
+    assert isinstance(finding, migration_data_checks.BlockedChange)
+    assert finding.rows == 1, "the check must count the row one byte over the line and not the row on it"
+    assert finding.detail == f"with an index entry over 2704 bytes, largest 2712 bytes ({described_as})"
+
+
+@requires_db
+def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None:
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute('CREATE TABLE "preflight_slow" ("a" INTEGER)')
+        conn.execute('INSERT INTO "preflight_slow" SELECT generate_series(1, 50)')
+        (check,) = migration_data_checks.checks_for_statement(
+            'ALTER TABLE "preflight_slow" ADD CONSTRAINT "slow" CHECK (pg_sleep(0.1) IS NULL)', "probe"
+        )
+        findings: Final = migration_data_checks.find_blockers(conn, "public", (check,), timeout_seconds=0.5)
+        timeout_afterwards: Final = conn.execute("SHOW statement_timeout").fetchone()
+
+    assert findings == (migration_data_checks.CouldNotVerify(check, 0.5),)
+    assert timeout_afterwards == ("0",), "the preflight must hand the session back with the timeout it found"
+    assert "could not be checked within 0.5 s" in migration_data_checks.report(findings)

tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py

@@ -0,0 +1,423 @@
+"""A Bedrock deployment whose ``guardrailConfig`` has ``trace: enabled`` bills the policy units Bedrock reports in the
+trace as guardrail cost, on every surface the deployment serves: ``/v1/chat/completions`` and ``/v1/responses`` over
+Converse and Invoke, streamed and not, and the four ``/bedrock/model/{model}`` pass-through actions. The cost rides
+``x-litellm-response-cost`` on non-streaming responses and the spend row's ``spend`` and ``guardrail_information``
+everywhere. ``trace: disabled`` leaves no usage, so nothing is billed beyond the tokens (LIT-10101)"""
+
+import base64
+import json
+import uuid
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from hashlib import sha256
+from pathlib import Path
+from typing import Final
+
+import boto3
+import pytest
+import yaml
+from botocore.config import Config
+from integration._support.budget import case_timeout_seconds
+from integration._support.client import Gateway, Scenario, eventually
+from integration._support.process import graceful_stop_seconds, owned_proxy
+from integration._support.upstream import _aws_event_frame
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+import litellm
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240)
+
+_MODEL_ID: Final = "anthropic.claude-sonnet-4-5-20250929-v1:0"
+_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
+_GUARDRAIL_ID: Final = "native-guardrail"
+_INPUT_COST_PER_TOKEN: Final = 0.001
+_OUTPUT_COST_PER_TOKEN: Final = 0.002
+_PROMPT_TOKENS: Final = 5
+_COMPLETION_TOKENS: Final = 1
+_INPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 1, "contentPolicyUnits": 1, "wordPolicyUnits": 1}
+_OUTPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 2, "sensitiveInformationPolicyUnits": 2}
+_SUMMED_USAGE: Final[Mapping[str, int]] = {
+    "topicPolicyUnits": 3,
+    "contentPolicyUnits": 1,
+    "wordPolicyUnits": 1,
+    "sensitiveInformationPolicyUnits": 2,
+}
+_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream")
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_SPEND_SQL: Final = 'SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"'
+
+
+def _llm_cost() -> float:
+    return _PROMPT_TOKENS * _INPUT_COST_PER_TOKEN + _COMPLETION_TOKENS * _OUTPUT_COST_PER_TOKEN
+
+
+def _guardrail_cost() -> float:
+    """What the proxy's own cost map prices the traced units at, so the test owns the relation, not AWS's rates."""
+    rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"]
+    return sum(units * rates[counter] for counter, units in _SUMMED_USAGE.items())
+
+
+def _action(request: Request) -> str:
+    return request.target.partition("?")[0].rsplit("/", 1)[-1]
+
+
+def _traced(request: Request) -> bool:
+    if _action(request).startswith("converse"):
+        config: Final = json.loads(request.body).get("guardrailConfig") or {}
+        return config.get("trace") == "enabled"
+    return request.headers.get("x-amzn-bedrock-trace", "").upper() == "ENABLED"
+
+
+def _assessment(usage: Mapping[str, int]) -> dict[str, JsonValue]:
+    return {"invocationMetrics": {"usage": dict(usage), "guardrailProcessingLatency": 7}}
+
+
+def _converse_trace() -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "inputAssessment": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)},
+            "outputAssessments": {_GUARDRAIL_ID: [_assessment(_OUTPUT_USAGE)]},
+        }
+    }
+
+
+def _invoke_trace() -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "input": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)},
+            "outputs": [{_GUARDRAIL_ID: _assessment(_OUTPUT_USAGE)}],
+        }
+    }
+
+
+def _bedrock(request: Request) -> Reply:
+    """Bedrock's runtime as the caller sees it: a one-token answer, and the guardrail trace with its policy units only
+    when the request enabled tracing, in the shape of the action that served it."""
+    traced: Final = _traced(request)
+    action: Final = _action(request)
+    usage: Final = {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6}
+    converse_trace: Final = {"trace": _converse_trace()} if traced else {}
+    if action == "converse":
+        body: Final = {
+            "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
+            "stopReason": "end_turn",
+            "usage": usage,
+            "metrics": {"latencyMs": 3},
+            **converse_trace,
+        }
+        return Reply(body=json.dumps(body).encode())
+    if action == "converse-stream" and "toolConfig" in json.loads(request.body):
+        return Reply(
+            chunks=(
+                _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"),
+                _aws_event_frame(
+                    "contentBlockStart",
+                    {"start": {"toolUse": {"toolUseId": "tool-1", "name": "lookup"}}, "contentBlockIndex": 0},
+                    "sc",
+                    "u",
+                ),
+                _aws_event_frame(
+                    "contentBlockDelta",
+                    {"delta": {"toolUse": {"input": '{"q": "x"}'}}, "contentBlockIndex": 0},
+                    "sc",
+                    "u",
+                ),
+                _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"),
+                _aws_event_frame("messageStop", {"stopReason": "tool_use"}, "sc", "u"),
+                _aws_event_frame(
+                    "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u"
+                ),
+            ),
+            content_type=_EVENT_STREAM,
+        )
+    if action == "converse-stream":
+        return Reply(
+            chunks=(
+                _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"),
+                _aws_event_frame("contentBlockDelta", {"delta": {"text": "ok"}, "contentBlockIndex": 0}, "sc", "u"),
+                _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"),
+                _aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"),
+                _aws_event_frame(
+                    "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u"
+                ),
+            ),
+            content_type=_EVENT_STREAM,
+        )
+    invoke_trace: Final = {"amazon-bedrock-trace": _invoke_trace()} if traced else {}
+    if action == "invoke":
+        message: Final = {
+            "id": f"msg_{uuid.uuid4().hex[:12]}",
+            "type": "message",
+            "role": "assistant",
+            "model": "claude",
+            "content": [{"type": "text", "text": "ok"}],
+            "stop_reason": "end_turn",
+            "stop_sequence": None,
+            "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": _COMPLETION_TOKENS},
+            **invoke_trace,
+        }
+        return Reply(body=json.dumps(message).encode())
+    events: Final = (
+        {
+            "type": "message_start",
+            "message": {
+                "id": f"msg_{uuid.uuid4().hex[:12]}",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude",
+                "content": [],
+                "stop_reason": None,
+                "stop_sequence": None,
+                "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": 0},
+            },
+        },
+        {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+        {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
+        {"type": "content_block_stop", "index": 0},
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": _COMPLETION_TOKENS},
+        },
+        {
+            "type": "message_stop",
+            "amazon-bedrock-invocationMetrics": {
+                "inputTokenCount": _PROMPT_TOKENS,
+                "outputTokenCount": _COMPLETION_TOKENS,
+                "invocationLatency": 3,
+                "firstByteLatency": 1,
+            },
+            **invoke_trace,
+        },
+    )
+    return Reply(
+        chunks=tuple(
+            _aws_event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}, "sc", "u")
+            for event in events
+        ),
+        content_type=_EVENT_STREAM,
+    )
+
+
+def _deployment(wire: Wire, model: str, trace: str) -> Mapping[str, JsonValue]:
+    return {
+        "model": model,
+        "api_base": wire.url,
+        "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
+        "aws_secret_access_key": "scripted-secret",
+        "aws_region_name": "us-east-1",
+        "use_in_pass_through": True,
+        "input_cost_per_token": _INPUT_COST_PER_TOKEN,
+        "output_cost_per_token": _OUTPUT_COST_PER_TOKEN,
+        "guardrailConfig": {"guardrailIdentifier": _GUARDRAIL_ID, "guardrailVersion": "1", "trace": trace},
+    }
+
+
+@dataclass(frozen=True, slots=True)
+class Rig:
+    proxy: Gateway
+    scenario: Scenario
+    converse_model: str
+    invoke_model: str
+    untraced_model: str
+
+    def key(self) -> str:
+        return self.scenario.key(key_alias=f"native-trace-{uuid.uuid4().hex[:8]}")
+
+    def runtime(self, monkeypatch: pytest.MonkeyPatch, key: str):  # noqa: ANN201  # boto3 clients are untyped
+        monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key)
+        return boto3.client(
+            "bedrock-runtime",
+            endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock",
+            region_name="us-east-1",
+            config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}),
+        )
+
+
+@pytest.fixture(scope="module")
+def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+    suffix: Final = uuid.uuid4().hex[:12]
+    converse_model: Final = f"bedrock-traced-converse-{suffix}"
+    invoke_model: Final = f"bedrock-traced-invoke-{suffix}"
+    untraced_model: Final = f"bedrock-untraced-{suffix}"
+    directory: Final = tmp_path_factory.mktemp("bedrock-native-guardrail-trace")
+    base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    with wire_server(_bedrock) as wire:
+        config: Final = {
+            **base_config,
+            "model_list": [
+                *base_config.get("model_list", []),
+                {"model_name": converse_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "enabled")},
+                {
+                    "model_name": invoke_model,
+                    "litellm_params": _deployment(wire, f"bedrock/invoke/{_MODEL_ID}", "enabled"),
+                },
+                {"model_name": untraced_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "disabled")},
+            ],
+        }
+        config_path: Final = directory / "bedrock-native-guardrail-trace.yaml"
+        config_path.write_text(yaml.safe_dump(config))
+        with owned_proxy(session_gateway, directory, {}, config=config_path) as proxy, proxy.scenario() as scenario:
+            yield Rig(
+                proxy=proxy,
+                scenario=scenario,
+                converse_model=converse_model,
+                invoke_model=invoke_model,
+                untraced_model=untraced_model,
+            )
+
+
+@dataclass(frozen=True, slots=True)
+class SpendRow:
+    spend: float
+    guardrail_information: tuple[dict[str, JsonValue], ...]
+
+
+def _spend_row(proxy: Gateway, key: str) -> SpendRow:
+    digest: Final = sha256(key.encode()).hexdigest()
+    (row,) = eventually(lambda: proxy.rows(_SPEND_SQL, (digest,)), lambda rows: len(rows) == 1, seconds=70)
+    metadata: Final = _JSON_OBJECT.validate_python(
+        json.loads(row["metadata"]) if isinstance(row["metadata"], str) else row["metadata"]
+    )
+    raw_information: Final = metadata.get("guardrail_information")
+    entries: Final = raw_information if isinstance(raw_information, list) else [raw_information]
+    return SpendRow(
+        spend=float(str(row["spend"])),
+        guardrail_information=tuple(_JSON_OBJECT.validate_python(entry) for entry in entries if entry is not None),
+    )
+
+
+def _assert_billed_as_guardrail_cost(row: SpendRow) -> None:
+    assert row.spend == pytest.approx(_llm_cost() + _guardrail_cost())
+    (entry,) = row.guardrail_information
+    assert entry["guardrail_provider"] == "bedrock"
+    assert entry["guardrail_name"] == _GUARDRAIL_ID
+    assert entry["guardrail_usage"] == _SUMMED_USAGE
+    assert isinstance(entry["guardrail_cost"], float) and entry["guardrail_cost"] == pytest.approx(_guardrail_cost())
+
+
+def _chat_body(model: str, stream: bool) -> dict[str, JsonValue]:
+    return {"model": model, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 16, "stream": stream}
+
+
+def _responses_body(model: str, stream: bool) -> dict[str, JsonValue]:
+    return {"model": model, "input": "hello", "stream": stream}
+
+
+@pytest.mark.parametrize("api", ["converse", "invoke"])
+@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"])
+def test_non_streaming_response_prices_the_trace_into_the_cost_header_and_the_spend_row(
+    rig: Rig, endpoint: str, api: str
+) -> None:
+    key: Final = rig.key()
+    model: Final = rig.converse_model if api == "converse" else rig.invoke_model
+    body: Final = _chat_body(model, False) if endpoint == "/v1/chat/completions" else _responses_body(model, False)
+    response: Final = rig.proxy.request("POST", endpoint, body, key=key)
+    assert response.is_success, response.text[:400]
+    assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost() + _guardrail_cost())
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+@pytest.mark.parametrize("api", ["converse", "invoke"])
+@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"])
+def test_streamed_response_prices_the_trace_of_the_final_event_into_the_spend_row(
+    rig: Rig, endpoint: str, api: str
+) -> None:
+    key: Final = rig.key()
+    model: Final = rig.converse_model if api == "converse" else rig.invoke_model
+    body: Final = _chat_body(model, True) if endpoint == "/v1/chat/completions" else _responses_body(model, True)
+    with rig.proxy.client.stream(
+        "POST", endpoint, json=body, headers={"Authorization": f"Bearer {key}"}, timeout=case_timeout_seconds()
+    ) as response:
+        frames: Final = tuple(line for line in response.iter_lines() if line.startswith("data:"))
+    assert response.is_success and frames
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+@pytest.mark.parametrize("action", _ACTIONS)
+def test_passthrough_action_prices_the_trace_into_the_spend_row(
+    rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str
+) -> None:
+    key: Final = rig.key()
+    runtime: Final = rig.runtime(monkeypatch, key)
+    messages: Final = [{"role": "user", "content": [{"text": "hello"}]}]
+    invoke_body: Final = json.dumps(
+        {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 16,
+            "messages": [{"role": "user", "content": "hello"}],
+        }
+    )
+    if action == "converse":
+        getattr(runtime, "converse")(modelId=rig.converse_model, messages=messages)
+    elif action == "converse-stream":
+        list(getattr(runtime, "converse_stream")(modelId=rig.converse_model, messages=messages)["stream"])
+    elif action == "invoke":
+        getattr(runtime, "invoke_model")(modelId=rig.converse_model, body=invoke_body)["body"].read()
+    else:
+        list(
+            getattr(runtime, "invoke_model_with_response_stream")(modelId=rig.converse_model, body=invoke_body)["body"]
+        )
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+def test_disabled_trace_bills_only_the_tokens(rig: Rig) -> None:
+    key: Final = rig.key()
+    response: Final = rig.proxy.request("POST", "/v1/chat/completions", _chat_body(rig.untraced_model, False), key=key)
+    assert response.is_success, response.text[:400]
+    assert "trace" not in response.json()
+    assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost())
+    row: Final = _spend_row(rig.proxy, key)
+    assert row.spend == pytest.approx(_llm_cost())
+    assert row.guardrail_information == ()
+
+
+def _messages_body(model: str, stream: bool) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": [{"role": "user", "content": f"hello from the messages api, stream={stream}"}],
+        "max_tokens": 16,
+        "stream": stream,
+    }
+
+
+def test_messages_non_streaming_response_applies_the_guardrail_and_prices_its_trace(rig: Rig) -> None:
+    key: Final = rig.key()
+    response: Final = rig.proxy.request("POST", "/v1/messages", _messages_body(rig.invoke_model, False), key=key)
+    assert response.is_success, response.text[:400]
+    assert "amazon-bedrock-trace" in response.json()
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+def test_messages_streamed_response_prices_the_trace_of_the_message_stop_event(rig: Rig) -> None:
+    key: Final = rig.key()
+    with rig.proxy.client.stream(
+        "POST",
+        "/v1/messages",
+        json=_messages_body(rig.invoke_model, True),
+        headers={"Authorization": f"Bearer {key}"},
+        timeout=case_timeout_seconds(),
+    ) as response:
+        body: Final = b"".join(response.iter_bytes())
+    assert response.is_success and b"message_stop" in body, body[:400]
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+def test_passthrough_streamed_tool_call_prices_the_trace_of_the_final_event(
+    rig: Rig, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    key: Final = rig.key()
+    runtime: Final = rig.runtime(monkeypatch, key)
+    tool_config: Final = {
+        "tools": [{"toolSpec": {"name": "lookup", "inputSchema": {"json": {"type": "object", "properties": {}}}}}]
+    }
+    events: Final = list(
+        getattr(runtime, "converse_stream")(
+            modelId=rig.converse_model,
+            messages=[{"role": "user", "content": [{"text": "hello"}]}],
+            toolConfig=tool_config,
+        )["stream"]
+    )
+    assert any("contentBlockStart" in event for event in events)
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))

tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py

@@ -1,7 +1,11 @@
+import json
+from typing import Final
 
 import pytest
+from pydantic import JsonValue
 
 import litellm
+from litellm.litellm_core_utils.llm_cost_calc import guardrail_cost
 from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
     bedrock_guardrail_cost,
     bedrock_guardrail_cost_by_unit,
@@ -230,3 +234,118 @@ def test_guardrail_information_cost_skips_malformed_entry_keeps_siblings():
     ]
     assert guardrail_information_cost(entries) == pytest.approx(0.0003)
     assert guardrail_information_cost({"guardrail_cost": 0.5, "guardrail_cost_in_spend": "maybe"}) == 0.0
+
+
+_TRACE_USAGE_IN: Final = {
+    "topicPolicyUnits": 1,
+    "contentPolicyUnits": 1,
+    "wordPolicyUnits": 1,
+    "sensitiveInformationPolicyUnits": 0,
+    "contextualGroundingPolicyUnits": 0,
+    "automatedReasoningPolicies": 0,
+}
+_TRACE_USAGE_OUT: Final = {
+    **_TRACE_USAGE_IN,
+    "topicPolicyUnits": 2,
+    "contentPolicyUnits": 2,
+    "sensitiveInformationPolicyUnits": 2,
+}
+_TRACE_PRICING: Final = {
+    "bedrock/guardrails": {
+        "litellm_provider": "bedrock",
+        "mode": "guardrails",
+        "guardrail_cost_per_unit": {
+            "topicPolicyUnits": 0.00015,
+            "contentPolicyUnits": 0.00015,
+            "wordPolicyUnits": 0.0,
+            "sensitiveInformationPolicyUnits": 0.0001,
+        },
+    }
+}
+
+
+def _converse_trace(action: str = "NONE") -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "inputAssessment": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}},
+            "outputAssessments": {
+                "gid": [
+                    {
+                        "topicPolicy": {"topics": [{"name": "t", "type": "DENY", "action": action}]},
+                        "invocationMetrics": {"usage": _TRACE_USAGE_OUT},
+                    }
+                ]
+            },
+        }
+    }
+
+
+def _invoke_trace() -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "input": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}},
+            "outputs": [{"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_OUT}}}],
+        }
+    }
+
+
+@pytest.mark.parametrize("trace", [_converse_trace(), _invoke_trace()], ids=["converse", "invoke"])
+def test_native_trace_sums_every_assessment_and_prices_it_like_the_proxy_guardrail(monkeypatch, trace) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    entry: Final = guardrail_cost.native_bedrock_guardrail_information(
+        trace, aws_region_name=None, guardrail_identifier=None
+    )
+    assert entry is not None
+    assert entry["guardrail_usage"] == {
+        "topicPolicyUnits": 3,
+        "contentPolicyUnits": 3,
+        "wordPolicyUnits": 2,
+        "sensitiveInformationPolicyUnits": 2,
+        "contextualGroundingPolicyUnits": 0,
+    }
+    assert entry["guardrail_cost"] == pytest.approx(3 * 0.00015 + 3 * 0.00015 + 2 * 0.0001)
+    assert entry["guardrail_cost_by_unit"]["contextualGroundingPolicyUnits"] is None
+    assert entry["guardrail_name"] == "gid"
+    assert entry["guardrail_provider"] == "bedrock"
+    assert entry["guardrail_status"] == "success"
+    assert guardrail_cost.guardrail_information_cost([entry]) == pytest.approx(entry["guardrail_cost"])
+
+
+def test_native_trace_marks_a_blocked_assessment_intervened_and_prefers_the_configured_identifier(monkeypatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    entry: Final = guardrail_cost.native_bedrock_guardrail_information(
+        _converse_trace("BLOCKED"), aws_region_name="us-east-1", guardrail_identifier="configured-id"
+    )
+    assert entry is not None
+    assert entry["guardrail_status"] == "guardrail_intervened"
+    assert entry["guardrail_name"] == "configured-id"
+
+
+def test_native_trace_without_usage_is_not_billed(monkeypatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    unpriced: Final = {"guardrail": {"inputAssessment": {"gid": {"topicPolicy": {"topics": []}}}}}
+    assert guardrail_cost.native_bedrock_guardrail_information(unpriced, None, None) is None
+    assert guardrail_cost.bedrock_guardrail_trace({"usage": {}}) is None
+
+
+def test_with_guardrail_entry_records_the_same_trace_once() -> None:
+    entry: Final = {"guardrail_name": "gid", "guardrail_cost": 0.001}
+    metadata: dict = {"standard_logging_guardrail_information": {"guardrail_name": "hook"}}
+    guardrail_cost.with_guardrail_entry(metadata, entry)  # pyright: ignore[reportArgumentType]  # partial entry
+    guardrail_cost.with_guardrail_entry(metadata, entry)  # pyright: ignore[reportArgumentType]  # partial entry
+    assert metadata["standard_logging_guardrail_information"] == [{"guardrail_name": "hook"}, entry]
+
+
+def test_native_trace_entry_redacts_the_pii_matches_before_they_reach_spend_logs(monkeypatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    trace: Final = _converse_trace()
+    trace["guardrail"]["inputAssessment"]["gid"]["sensitiveInformationPolicy"] = {
+        "piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}]
+    }
+    entry: Final = guardrail_cost.native_bedrock_guardrail_information(
+        trace, aws_region_name=None, guardrail_identifier=None
+    )
+    assert entry is not None
+    assert entry["guardrail_status"] == "guardrail_intervened"
+    assert "123-45-6789" not in json.dumps(entry["guardrail_response"])
+    assert entry["guardrail_usage"]["sensitiveInformationPolicyUnits"] == 2

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -3672,6 +3672,73 @@ def test_get_extra_header_tags():
                 delattr(litellm, "extra_spend_tag_headers")
 
 
+def test_get_extra_header_tags_matches_title_case_config_with_lowercase_request_header(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "extra_spend_tag_headers", ["X-Org-Id"], raising=False)
+
+    assert StandardLoggingPayloadSetup._get_extra_header_tags(
+        proxy_server_request={"headers": {"x-org-id": "org_1"}}
+    ) == ["X-Org-Id: org_1"]
+
+
+def test_get_extra_header_tags_matches_mixed_case_request_header_with_lowercase_config(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "extra_spend_tag_headers", ["x-customer-id"], raising=False)
+
+    assert StandardLoggingPayloadSetup._get_extra_header_tags(
+        proxy_server_request={"headers": {"X-Customer-Id": "v"}}
+    ) == ["x-customer-id: v"]
+
+
+def test_get_extra_header_tags_keeps_lowercase_label(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "extra_spend_tag_headers", ["x-customer-id"], raising=False)
+
+    assert StandardLoggingPayloadSetup._get_extra_header_tags(
+        proxy_server_request={"headers": {"x-customer-id": "v"}}
+    ) == ["x-customer-id: v"]
+
+
+def test_get_extra_header_tags_returns_none_for_absent_or_empty_header(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "extra_spend_tag_headers", ["x-org-id"], raising=False)
+
+    assert StandardLoggingPayloadSetup._get_extra_header_tags(proxy_server_request={"headers": {}}) is None
+    assert StandardLoggingPayloadSetup._get_extra_header_tags(
+        proxy_server_request={"headers": {"x-org-id": ""}}
+    ) is None
+
+
+def test_get_extra_header_tags_preserves_configured_header_order(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "extra_spend_tag_headers", ["x-second", "x-first"], raising=False)
+
+    assert StandardLoggingPayloadSetup._get_extra_header_tags(
+        proxy_server_request={"headers": {"x-first": "1", "x-second": "2"}}
+    ) == ["x-second: 2", "x-first: 1"]
+
+
+def test_get_request_tags_includes_extra_header_tag_with_user_agent_tags(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "extra_spend_tag_headers", ["X-Org-Id"], raising=False)
+
+    assert StandardLoggingPayloadSetup.get_request_tags(
+        litellm_params={},
+        proxy_server_request={
+            "headers": {
+                "user-agent": "litellm/1.0",
+                "x-org-id": "org_1",
+            }
+        },
+    ) == ["User-Agent: litellm", "User-Agent: litellm/1.0", "X-Org-Id: org_1"]
+
+
 def test_response_cost_calculator_with_response_cost_in_hidden_params(logging_obj):
     from litellm import Router
 
@@ -13642,3 +13709,54 @@ def test_native_messages_body_carries_the_reported_list_cost_into_the_priced_res
     assert isinstance(priced, ModelResponse)
     headers = priced._hidden_params.get("additional_headers") or {}
     assert headers.get("llm_provider-x-litellm-response-cost") == expected_header
+
+
+def test_record_native_bedrock_guardrail_bills_when_the_metadata_is_none() -> None:
+    import datetime
+
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    logging_obj = Logging(
+        model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
+        messages=[{"role": "user", "content": "hello"}],
+        stream=False,
+        call_type="completion",
+        start_time=datetime.datetime.now(),
+        litellm_call_id="call-none-metadata",
+        function_id="fn",
+    )
+    logging_obj.model_call_details["litellm_params"] = {"metadata": None}
+    trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}}
+
+    logging_obj.record_native_bedrock_guardrail({"trace": trace})
+
+    (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
+    assert entry["guardrail_usage"] == {"contentPolicyUnits": 2}
+    assert entry["guardrail_cost"] > 0
+
+
+def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_from_cache() -> None:
+    import datetime
+
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    logging_obj = Logging(
+        model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
+        messages=[{"role": "user", "content": "hello"}],
+        stream=False,
+        call_type="completion",
+        start_time=datetime.datetime.now(),
+        litellm_call_id="call-cache-hit",
+        function_id="fn",
+    )
+    logging_obj.model_call_details["litellm_params"] = {"metadata": {}}
+    trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}}
+    cached = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x")
+    setattr(cached, "trace", trace)
+    cached._hidden_params = {"cache_hit": True}
+
+    logging_obj.record_native_bedrock_guardrail(cached)
+    logging_obj.model_call_details["cache_hit"] = True
+    logging_obj.record_native_bedrock_guardrail({"trace": trace})
+
+    assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"]

tests/unit/litellm_core_utils/test_redact_messages.py

@@ -26,15 +26,16 @@ import litellm
 from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.litellm_core_utils.redact_messages import (
     _redact_responses_api_output,
     _redacted_responses_api_response,
     perform_redaction,
     redacted_model_call_details_for_custom_logger,
+    redacted_response_obj_for_custom_logger,
     redacted_standard_logging_payload,
     should_redact_message_logging,
 )
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.responses.main import mock_responses_api_response
 from litellm.types.router import BaselineRouteStamp
 from litellm.types.utils import (
@@ -176,7 +177,11 @@ def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str
         object="response",
         status="completed",
         output=[
-            {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": SECRET})}
+            {
+                "type": "compaction",
+                "id": "cmp_1",
+                "encrypted_content": json.dumps({"type": "compaction", "content": SECRET}),
+            }
         ],
     )
     result: Final = response if shape == "object" else response.model_dump()
@@ -2041,3 +2046,32 @@ async def test_redaction_with_metadata_completion_api():
     response = standard_logging_payload["response"]
     assert response["choices"][0]["message"]["content"] == "redacted-by-litellm"
     assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
+
+
+def test_perform_redaction_strips_the_matched_text_from_the_bedrock_guardrail_traces() -> None:
+    finding = {
+        "guardrail": {
+            "inputAssessment": {
+                "gid": {
+                    "sensitiveInformationPolicy": {
+                        "piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}]
+                    }
+                }
+            }
+        }
+    }
+    response = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x")
+    setattr(response, "trace", copy.deepcopy(finding))
+    setattr(response, "amazon-bedrock-trace", copy.deepcopy(finding))
+    opted_out: Final = CustomLogger(message_logging=False)
+
+    redacted: Final = redacted_response_obj_for_custom_logger(response, opted_out)
+    redacted_dict: Final = redacted_response_obj_for_custom_logger(
+        {"choices": [], "trace": copy.deepcopy(finding)}, opted_out
+    )
+
+    assert isinstance(redacted, litellm.ModelResponse)
+    assert "123-45-6789" not in redacted.model_dump_json()
+    assert "US_SSN" in redacted.model_dump_json(), "the finding stays, only its matched text goes"
+    assert "123-45-6789" not in str(redacted_dict)
+    assert "123-45-6789" in response.model_dump_json(), "the caller's own response is not redacted"

tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -8,17 +8,19 @@ import httpx
 import pytest
 import respx
 
-
 import litellm
 from litellm import ChatCompletionUsageBlock, stream_chunk_builder
-from litellm.types.utils import GenericStreamingChunk
-from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor, stream_chunk_builder_text_completion
+from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
+    ChunkProcessor,
+    stream_chunk_builder_text_completion,
+)
 from litellm.llms.anthropic.chat.handler import ModelResponseIterator
 from litellm.types.utils import (
     ChatCompletionDeltaToolCall,
     ChatCompletionMessageToolCall,
     Delta,
     Function,
+    GenericStreamingChunk,
     ModelResponseStream,
     PromptTokensDetails,
     ServerToolUse,
@@ -196,11 +198,7 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
                 }
             ]
         ),
-        make_chunk(
-            thinking_blocks=[
-                {"type": "thinking", "thinking": None, "signature": "sig_block1"}
-            ]
-        ),
+        make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block1"}]),
         make_chunk(
             thinking_blocks=[
                 {
@@ -218,16 +216,10 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
                 }
             ]
         ),
-        make_chunk(
-            thinking_blocks=[
-                {"type": "thinking", "thinking": None, "signature": "sig_block2"}
-            ]
-        ),
+        make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block2"}]),
     ]
 
-    thinking_chunks = [
-        chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks")
-    ]
+    thinking_chunks = [chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks")]
     processor = ChunkProcessor(chunks=chunks)
     result = processor.get_combined_thinking_content(thinking_chunks)
 
@@ -297,9 +289,7 @@ def test_cache_read_input_tokens_retained():
             prompt_tokens=11779,
             total_tokens=11784,
             completion_tokens_details=None,
-            prompt_tokens_details=PromptTokensDetails(
-                audio_tokens=None, cached_tokens=11775
-            ),
+            prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=11775),
             cache_creation_input_tokens=4,
             cache_read_input_tokens=11775,
         ),
@@ -333,9 +323,7 @@ def test_cache_read_input_tokens_retained():
             prompt_tokens=0,
             total_tokens=214,
             completion_tokens_details=None,
-            prompt_tokens_details=PromptTokensDetails(
-                audio_tokens=None, cached_tokens=0
-            ),
+            prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=0),
             cache_creation_input_tokens=0,
             cache_read_input_tokens=0,
         ),
@@ -395,10 +383,7 @@ def test_streaming_preserves_anthropic_1hr_cache_creation_breakdown():
     )
     # Sanity: the delta event genuinely lacks the breakdown - this is the input
     # condition that used to defeat cost calc.
-    assert (
-        getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None)
-        is None
-    )
+    assert getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None) is None
 
     def _usage_chunk(usage, finish_reason):
         return ModelResponseStream(
@@ -517,9 +502,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk():
             prompt_tokens=1234,
             total_tokens=1239,
             completion_tokens_details=None,
-            prompt_tokens_details=PromptTokensDetails(
-                audio_tokens=None, cached_tokens=543
-            ).model_dump(),
+            prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=543).model_dump(),
         ),
         index=2,
     )
@@ -536,6 +519,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk():
 
     assert usage.prompt_tokens_details.cached_tokens == 543
 
+
 def test_stream_chunk_builder_litellm_usage_chunks():
     """
     Validate ChunkProcessor.calculate_usage uses provided usage fields from streaming chunks
@@ -609,9 +593,7 @@ def test_stream_chunk_builder_litellm_usage_chunks():
     chunks = [chunk1, chunk2]
     processor = ChunkProcessor(chunks=chunks)
 
-    usage = processor.calculate_usage(
-        chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output=""
-    )
+    usage = processor.calculate_usage(chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output="")
 
     assert usage.prompt_tokens == 50
     assert usage.completion_tokens == 27
@@ -655,15 +637,11 @@ def test_calculate_usage_honors_openai_sdk_completion_usage_chunks():
         provider_specific_fields=None,
         stream_options={"include_usage": True},
     )
-    usage_chunk.usage = CompletionUsage(
-        prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704
-    )
+    usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704)
     assert type(usage_chunk.usage) is CompletionUsage
 
     chunks = [content_chunk, usage_chunk]
-    usage = ChunkProcessor(chunks=chunks).calculate_usage(
-        chunks=chunks, model="mantle-claude", completion_output=""
-    )
+    usage = ChunkProcessor(chunks=chunks).calculate_usage(chunks=chunks, model="mantle-claude", completion_output="")
 
     assert usage.prompt_tokens == 20
     assert usage.completion_tokens == 60
@@ -686,9 +664,7 @@ def test_get_model_from_chunks_azure_model_router():
         {"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []},
     ]
 
-    result = ChunkProcessor._get_model_from_chunks(
-        chunks=chunks, first_chunk_model="azure-model-router"
-    )
+    result = ChunkProcessor._get_model_from_chunks(chunks=chunks, first_chunk_model="azure-model-router")
 
     # Should return the actual model, not the request model
     assert result == "gpt-4.1-nano-2025-04-14"
@@ -699,9 +675,7 @@ def test_get_model_from_chunks_azure_model_router():
         {"model": "gpt-4", "id": "chatcmpl-456", "choices": []},
     ]
 
-    result_same = ChunkProcessor._get_model_from_chunks(
-        chunks=chunks_same_model, first_chunk_model="gpt-4"
-    )
+    result_same = ChunkProcessor._get_model_from_chunks(chunks=chunks_same_model, first_chunk_model="gpt-4")
 
     # Should return the first chunk's model when all are the same
     assert result_same == "gpt-4"
@@ -777,9 +751,7 @@ def test_stream_chunk_builder_anthropic_web_search():
     chunks = [chunk1, chunk2]
     processor = ChunkProcessor(chunks=chunks)
 
-    usage = processor.calculate_usage(
-        chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output=""
-    )
+    usage = processor.calculate_usage(chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output="")
 
     assert usage.prompt_tokens == 50
     assert usage.completion_tokens == 27
@@ -931,15 +903,11 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields():
         ],
     )
     chunk_dict = chunk.model_dump()
-    chunk_dict["_hidden_params"] = {
-        "provider_specific_fields": {"traffic_type": "default"}
-    }
+    chunk_dict["_hidden_params"] = {"provider_specific_fields": {"traffic_type": "default"}}
 
     response = stream_chunk_builder(chunks=[chunk_dict])
     assert response is not None
-    assert (
-        response._hidden_params["provider_specific_fields"]["traffic_type"] == "default"
-    )
+    assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default"
 
 
 def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks():
@@ -984,10 +952,7 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks():
     assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata
     assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata
     assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata
-    assert (
-        response._hidden_params["vertex_ai_url_context_metadata"]
-        == url_context_metadata
-    )
+    assert response._hidden_params["vertex_ai_url_context_metadata"] == url_context_metadata
 
     dumped = response.model_dump()
     assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata
@@ -1034,9 +999,7 @@ def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata():
 
 def test_stream_chunk_builder_propagates_vertex_ai_safety_results():
     """Assembled response must expose safety data under the non-streaming field name."""
-    safety_ratings = [
-        [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}]
-    ]
+    safety_ratings = [[{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}]]
 
     chunk = ModelResponseStream(
         id="chatcmpl-vertex-safety",
@@ -1078,18 +1041,12 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks():
             )
         ],
     ).model_dump()
-    chunk_dict["_hidden_params"] = {
-        "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}]
-    }
+    chunk_dict["_hidden_params"] = {"vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}]}
 
     response = stream_chunk_builder(chunks=[chunk_dict])
     assert response is not None
-    assert getattr(response, "vertex_ai_grounding_metadata") == [
-        {"webSearchQueries": ["test query"]}
-    ]
-    assert response.model_dump()["vertex_ai_grounding_metadata"] == [
-        {"webSearchQueries": ["test query"]}
-    ]
+    assert getattr(response, "vertex_ai_grounding_metadata") == [{"webSearchQueries": ["test query"]}]
+    assert response.model_dump()["vertex_ai_grounding_metadata"] == [{"webSearchQueries": ["test query"]}]
 
 
 def test_cost_field_in_usage_chunks():
@@ -1098,29 +1055,21 @@ def test_cost_field_in_usage_chunks():
         id="chatcmpl-1",
         created=1745513206,
         model="openrouter/claude",
-        choices=[
-            StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
-        ],
+        choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))],
         usage=chunk1_usage,
     )
 
-    chunk2_usage = Usage(
-        completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
-    )
+    chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025)
     chunk2 = ModelResponseStream(
         id="chatcmpl-1",
         created=1745513207,
         model="openrouter/claude",
-        choices=[
-            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
-        ],
+        choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))],
         usage=chunk2_usage,
     )
 
     processor = ChunkProcessor(chunks=[chunk1, chunk2])
-    usage = processor.calculate_usage(
-        chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi"
-    )
+    usage = processor.calculate_usage(chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi")
 
     assert hasattr(usage, "cost")
     assert usage.cost == 0.00025
@@ -1203,25 +1152,19 @@ def test_prompt_tokens_details_survive_later_usage_chunk_without_details():
         id="chatcmpl-1",
         created=1745513206,
         model="openai/gpt-5.6-sol",
-        choices=[
-            StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
-        ],
+        choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))],
         usage=Usage(
             prompt_tokens=6017,
             completion_tokens=4,
             total_tokens=6021,
-            prompt_tokens_details=PromptTokensDetailsWrapper(
-                cached_tokens=6004, cache_write_tokens=10
-            ),
+            prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=6004, cache_write_tokens=10),
         ),
     )
     chunk_without_details = ModelResponseStream(
         id="chatcmpl-1",
         created=1745513207,
         model="openai/gpt-5.6-sol",
-        choices=[
-            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
-        ],
+        choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))],
         usage=Usage(prompt_tokens=6017, completion_tokens=4, total_tokens=6021),
     )
 
@@ -1559,9 +1502,7 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate(
     assert usage.completion_tokens_details.text_tokens == expected_text_tokens
 
 
-def _openai_chunk(
-    choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None
-) -> dict[str, object]:
+def _openai_chunk(choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None) -> dict[str, object]:
     base: Final = {
         "id": "chatcmpl-lit6552",
         "object": "chat.completion.chunk",
@@ -1996,7 +1937,11 @@ def test_stream_chunk_builder_litellm_tool_call(respx_mock: respx.MockRouter) ->
     )
 
     assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [
-        {"function": {"arguments": '{"city":"Boston"}', "name": "get_weather"}, "id": "call_weather", "type": "function"}
+        {
+            "function": {"arguments": '{"city":"Boston"}', "name": "get_weather"},
+            "id": "call_weather",
+            "type": "function",
+        }
     ]
     assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (5, 3, 8)
 
@@ -2018,7 +1963,9 @@ def test_stream_chunk_builder_litellm_tool_call_regular_message(respx_mock: resp
 
 
 def test_stream_chunk_builder_multiple_tool_calls() -> None:
-    def tool_call_chunk(*tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None) -> ModelResponseStream:
+    def tool_call_chunk(
+        *tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None
+    ) -> ModelResponseStream:
         return _legacy_stream_chunk("openai/gpt-4o", tool_calls=tool_calls, finish_reason=finish_reason)
 
     chunks: Final = (
@@ -2029,7 +1976,9 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None:
         ),
         tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='{"base": 3, '))),
         tool_call_chunk(
-            ChatCompletionDeltaToolCall(id="call_add", index=1, type="function", function=Function(name="add", arguments=""))
+            ChatCompletionDeltaToolCall(
+                id="call_add", index=1, type="function", function=Function(name="add", arguments="")
+            )
         ),
         tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"exponent": 5}'))),
         tool_call_chunk(
@@ -2043,7 +1992,11 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None:
     assert response is not None
     assert response.choices[0].finish_reason == "tool_calls"
     assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [
-        {"function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"}, "id": "call_exp", "type": "function"},
+        {
+            "function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"},
+            "id": "call_exp",
+            "type": "function",
+        },
         {"function": {"arguments": '{"x": 1, "y": 2}', "name": "add"}, "id": "call_add", "type": "function"},
     ]
 
@@ -2077,9 +2030,7 @@ def test_grok_bug() -> None:
     response: Final = stream_chunk_builder(chunks=chunks)
 
     assert response is not None
-    assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {
-        "location": "Boston"
-    }
+    assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {"location": "Boston"}
 
 
 def test_stream_chunk_builder_openai_audio_output_usage() -> None:
@@ -2196,7 +2147,9 @@ def test_stream_chunk_builder_keeps_usage_dict_set_on_a_typed_text_completion_ch
 
 
 @pytest.mark.parametrize("build", [litellm.stream_chunk_builder, litellm.stream_chunk_builder_text_completion])
-def test_stream_chunk_builder_rebuilds_dict_text_completion_chunks(build: Callable[..., TextCompletionResponse]) -> None:
+def test_stream_chunk_builder_rebuilds_dict_text_completion_chunks(
+    build: Callable[..., TextCompletionResponse],
+) -> None:
     def chunk(choices: list[dict[str, object]], **extra: object) -> dict[str, object]:
         return {
             "id": "cmpl-1",
@@ -2450,7 +2403,9 @@ def test_stream_chunk_builder_rebuilds_text_completion_choice_indexes_that_are_n
         ]
     )
 
-    assert tuple(choice.model_dump(include={"index", "text", "finish_reason"}) for choice in response.choices) == expected
+    assert (
+        tuple(choice.model_dump(include={"index", "text", "finish_reason"}) for choice in response.choices) == expected
+    )
 
 
 def test_stream_chunk_builder_text_completion_survives_uncountable_prompt() -> None:
@@ -2517,3 +2472,50 @@ def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies():
     assert response is not None
     assert response.choices[0].message.content == "Hello world"
     assert response.choices[0].message.reasoning_items == reasoning_items
+
+
+def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it_on_the_logging_object(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A Bedrock stream reports its guardrail trace on the last chunk only; the built response carries it like the
+    non-streaming response and the call bills it as guardrail cost."""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+    monkeypatch.setattr(
+        litellm,
+        "model_cost",
+        {
+            "bedrock/guardrails": {
+                "litellm_provider": "bedrock",
+                "mode": "guardrails",
+                "guardrail_cost_per_unit": {"topicPolicyUnits": 0.00015},
+            }
+        },
+    )
+    trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 2}}}}}}
+    chunks = [
+        ModelResponseStream(id="s", model="m", choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]),
+        ModelResponseStream(
+            id="s",
+            model="m",
+            choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")],
+            provider_specific_fields={"trace": trace},
+        ),
+    ]
+    logging_obj = Logging(
+        model="bedrock/m",
+        messages=[{"role": "user", "content": "q"}],
+        stream=True,
+        call_type="acompletion",
+        start_time=None,
+        litellm_call_id="call",
+        function_id="fn",
+    )
+    logging_obj.model_call_details["litellm_params"] = {"metadata": {}}
+    built = stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "q"}], logging_obj=logging_obj)
+    assert built is not None
+    assert getattr(built, "trace") == trace
+    (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
+    assert entry["guardrail_usage"] == {"topicPolicyUnits": 2}
+    assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015)

tests/unit/litellm_proxy_extras/test_migration_data_checks.py

@@ -0,0 +1,156 @@
+from pathlib import Path
+from typing import Final
+
+import pytest
+from litellm_proxy_extras.migration_data_checks import (
+    DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS,
+    MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR,
+    BlockedChange,
+    Check,
+    CouldNotVerify,
+    btree_max_entry_bytes,
+    checks_for_index,
+    checks_for_migration,
+    checks_for_statement,
+    migration_data_check_timeout_seconds,
+    report,
+    statements,
+)
+
+MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
+API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0: Final = 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime")'
+
+
+def _shape(checks: tuple[Check, ...]) -> tuple[tuple[str, str, tuple[str, ...]], ...]:
+    return tuple((check.kind, check.table, check.columns) for check in checks)
+
+
+@pytest.mark.parametrize(
+    ("statement", "expected"),
+    (
+        (
+            API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0,
+            (("too_long", "LiteLLM_SpendLogs", ("api_key", "startTime")),),
+        ),
+        (
+            'CREATE UNIQUE INDEX CONCURRENTLY "T_a_b_key" ON "public"."T" USING btree ("a", "b" DESC)',
+            (("too_long", "T", ("a", "b")), ("duplicates", "T", ("a", "b"))),
+        ),
+        ('CREATE INDEX "T_a_idx" ON "T" USING gin ("a")', ()),
+        ('CREATE INDEX "T_a_idx" ON "T" (lower("a"))', ()),
+        ('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', (("nulls", "T", ("a",)),)),
+        ('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL', (("any_row", "T", ("a",)),)),
+        ('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL DEFAULT \'x\'', ()),
+        ('ALTER TABLE "T" ADD COLUMN "a" TEXT', ()),
+        (
+            'ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") ON DELETE CASCADE',
+            (("orphans", "T", ("r",)),),
+        ),
+        ('ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") NOT VALID', ()),
+        ('ALTER TABLE "T" ADD CONSTRAINT "T_a_key" UNIQUE ("a", "b")', (("duplicates", "T", ("a", "b")),)),
+        (
+            'ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0 AND "b" IS NOT NULL)',
+            (("check_violations", "T", ()),),
+        ),
+        ('ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0) NOT VALID', ()),
+        (
+            'ALTER TABLE "T" ADD COLUMN "c" TEXT, ALTER COLUMN "b" SET NOT NULL, ADD CONSTRAINT "u" UNIQUE ("c")',
+            (("nulls", "T", ("b",)), ("duplicates", "T", ("c",))),
+        ),
+        ('CREATE TABLE "T" ("id" TEXT NOT NULL, CONSTRAINT "T_pkey" PRIMARY KEY ("id"))', ()),
+        ('DROP INDEX IF EXISTS "T_a_idx"', ()),
+        ("SELECT 1", ()),
+    ),
+)
+def test_a_statement_gets_one_check_per_way_existing_rows_can_make_it_fail(
+    statement: str, expected: tuple[tuple[str, str, tuple[str, ...]], ...]
+) -> None:
+    assert _shape(checks_for_statement(statement, "migration x")) == expected
+
+
+def test_a_foreign_key_check_keeps_the_referenced_table_and_columns() -> None:
+    (check,) = checks_for_statement(
+        'ALTER TABLE "T" ADD CONSTRAINT "fk" FOREIGN KEY ("r", "s") REFERENCES "R"("id", "sub")', "migration x"
+    )
+    assert (check.referenced_table, check.referenced_columns) == ("R", ("id", "sub"))
+
+
+def test_a_check_constraint_keeps_its_expression_for_the_count_query() -> None:
+    (check,) = checks_for_statement('ALTER TABLE "T" ADD CONSTRAINT "c" CHECK ("a" > 0)', "migration x")
+    assert check.expression == '"a" > 0'
+
+
+def test_comments_and_blank_statements_are_dropped_before_matching() -> None:
+    migration_sql: Final = "-- CreateIndex\n" + API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0 + ";\n\n-- done; really\n"
+    assert statements(migration_sql) == (API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0,)
+    assert _shape(checks_for_migration("x", migration_sql)) == _shape(
+        checks_for_statement(API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0, "x")
+    )
+
+
+def test_every_shipped_migration_is_readable_and_the_later_indexes_on_text_columns_get_a_check() -> None:
+    derived: Final = {
+        folder.name: checks_for_migration(folder.name, (folder / "migration.sql").read_text())
+        for folder in MIGRATIONS.iterdir()
+        if folder.is_dir()
+    }
+    assert len(derived) > 200
+    everywhere: Final = tuple(check for checks in derived.values() for check in checks)  # comprehension-ok: flattening
+    assert any(
+        check.kind == "too_long" and check.table == "LiteLLM_SpendLogs" and "session_id" in check.columns
+        for check in everywhere
+    )
+
+
+def test_a_request_log_index_gets_one_entry_size_check_over_all_its_columns() -> None:
+    assert _shape(
+        checks_for_index("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime"))
+    ) == (("too_long", "LiteLLM_SpendLogs", ("api_key", "startTime")),)
+
+
+def test_the_btree_entry_limit_on_the_default_page_is_the_number_in_postgres_error() -> None:
+    # "index row size 2976 exceeds btree version 4 maximum 2704", Postgres 14, 8 KB page (BTMaxItemSize, nbtree.h)
+    assert btree_max_entry_bytes(8192) == 2704
+
+
+def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_was_applied() -> None:
+    (check,) = checks_for_index(
+        "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")
+    )
+    text: Final = report(
+        (BlockedChange(check, 2, "with an index entry over 2704 bytes, largest 3232 bytes (api_key up to 3207 bytes)"),)
+    )
+    assert (
+        '"LiteLLM_SpendLogs".api_key, startTime: 2 row(s) with an index entry over 2704 bytes, largest 3232 bytes '
+        "(api_key up to 3207 bytes)"
+    ) in text
+    assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text
+    assert 'UPDATE "LiteLLM_SpendLogs" SET api_key = encode(sha256(' in text
+    assert text.startswith("Existing rows would make a pending schema change fail:")
+    assert "No schema change was made" in text
+
+
+def test_a_scan_that_timed_out_is_reported_as_unverified_with_the_knob_to_turn() -> None:
+    (check,) = checks_for_statement('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', "migration 20270101000000_x")
+    text: Final = report((CouldNotVerify(check, 600.0),))
+    assert '"T".a: could not be checked within 600 s, refusing rather than guessing' in text
+    assert 'blocks: ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL (migration 20270101000000_x)' in text
+    assert MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR in text
+
+
+@pytest.mark.parametrize(
+    ("raw", "expected"),
+    (
+        (None, DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS),
+        ("30", 30.0),
+        ("nope", DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS),
+    ),
+)
+def test_the_scan_timeout_comes_from_the_environment_and_falls_back_on_garbage(
+    monkeypatch: pytest.MonkeyPatch, raw: str | None, expected: float
+) -> None:
+    if raw is None:
+        monkeypatch.delenv(MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR, raising=False)
+    else:
+        monkeypatch.setenv(MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR, raw)
+    assert migration_data_check_timeout_seconds() == expected

tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py

@@ -7,9 +7,12 @@ to be keyed on that same value rather than on the provider's own completion id.
 
 import datetime
 import json
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
+from typing import Final
 
 import pytest
 import respx
+from pydantic import TypeAdapter
 
 import litellm
 from litellm.llms.anthropic.messages.adapters.handler import (
@@ -18,6 +21,153 @@ from litellm.llms.anthropic.messages.adapters.handler import (
 from litellm.llms.anthropic.messages.adapters.streaming_iterator import (
     AnthropicStreamWrapper,
 )
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
+from litellm.types.utils import ModelResponseStream
+
+
+_EVENT: Final = TypeAdapter(Mapping[str, object])
+_MESSAGES: Final = TypeAdapter(list[AllMessageValues])
+_TOOLS: Final = TypeAdapter(list[ChatCompletionToolParam])
+_BLOCK_PROMPTS: Final[tuple[list[AllMessageValues], ...]] = (
+    [{"role": "user", "content": [{"type": "text", "text": "Count every word. " * 100}]}],
+    [{"role": "tool", "content": [{"type": "text", "text": "Count every word. " * 100}], "tool_call_id": "call_1"}],
+    [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": "https://image.example/test.png", "detail": "high"}}
+            ],
+        }
+    ],
+)
+
+
+def _assert_estimated_start_and_provider_final_usage(raw_events: Sequence[bytes], request_body: bytes) -> None:
+    request: Final = _EVENT.validate_json(request_body)
+    events: Final = tuple(_EVENT.validate_json(chunk.decode().split("data: ", 1)[1]) for chunk in raw_events)
+    start: Final = _EVENT.validate_python(
+        next(event for event in events if event["type"] == "message_start")["message"]
+    )
+    final: Final = next(event for event in events if event["type"] == "message_delta")
+    expected: Final = litellm.token_counter(
+        model="kimi-k2",
+        messages=_MESSAGES.validate_python(request["messages"]),
+        tools=_TOOLS.validate_python(request["tools"]),
+    )
+    assert expected > 0
+    assert start["usage"] == {
+        "input_tokens": expected,
+        "output_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        "cache_read_input_tokens": 0,
+    }
+    assert final["usage"] == {"input_tokens": 3, "output_tokens": 4}
+
+
+def test_sync_message_start_estimates_prompt_without_replacing_provider_usage(respx_mock: respx.MockRouter) -> None:
+    stream: Final = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(
+        max_tokens=16,
+        messages=MESSAGES,
+        system="Reply with a short greeting.",
+        tools=[{"name": "greet", "description": "Greet a person", "input_schema": {"type": "object"}}],
+        model="groq/kimi-k2",
+        stream=True,
+        custom_llm_provider="groq",
+    )
+    _assert_estimated_start_and_provider_final_usage(list(stream), respx_mock.calls.last.request.content)
+
+
+@pytest.mark.asyncio
+async def test_async_message_start_estimates_prompt_without_replacing_provider_usage(
+    respx_mock: respx.MockRouter,
+) -> None:
+    stream: Final = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler(
+        max_tokens=16,
+        messages=MESSAGES,
+        system="Reply with a short greeting.",
+        tools=[{"name": "greet", "description": "Greet a person", "input_schema": {"type": "object"}}],
+        model="groq/kimi-k2",
+        stream=True,
+        custom_llm_provider="groq",
+    )
+    raw: Final = [chunk async for chunk in stream]
+    _assert_estimated_start_and_provider_final_usage(raw, respx_mock.calls.last.request.content)
+
+
+class _FailingUpstream:
+    def __iter__(self) -> Iterator[ModelResponseStream]:
+        raise RuntimeError("upstream failed before its first chunk")
+
+    def __aiter__(self) -> AsyncIterator[ModelResponseStream]:
+        raise RuntimeError("upstream failed before its first chunk")
+
+
+@pytest.mark.parametrize("messages", _BLOCK_PROMPTS)
+def test_sync_initial_estimate_counts_content_blocks(
+    messages: list[AllMessageValues], respx_mock: respx.MockRouter
+) -> None:
+    image: Final = respx_mock.get("https://image.example/test.png").respond(503)
+    wrapper: Final = AnthropicStreamWrapper(iter([]), model="gpt-4o", completion_request={"messages": messages})
+    expected: Final = litellm.token_counter(model="gpt-4o", messages=messages, use_default_image_token_count=True)
+    first: Final = _EVENT.validate_python(next(wrapper))
+    message: Final = _EVENT.validate_python(first["message"])
+    assert message["usage"] == {
+        "input_tokens": expected,
+        "output_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        "cache_read_input_tokens": 0,
+    }
+    assert not image.called
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("messages", _BLOCK_PROMPTS)
+async def test_async_initial_estimate_counts_content_blocks(
+    messages: list[AllMessageValues], respx_mock: respx.MockRouter
+) -> None:
+    image: Final = respx_mock.get("https://image.example/test.png").respond(503)
+    wrapper: Final = AnthropicStreamWrapper(iter([]), model="gpt-4o", completion_request={"messages": messages})
+    expected: Final = litellm.token_counter(model="gpt-4o", messages=messages, use_default_image_token_count=True)
+    first: Final = _EVENT.validate_python(await anext(wrapper))
+    message: Final = _EVENT.validate_python(first["message"])
+    assert message["usage"] == {
+        "input_tokens": expected,
+        "output_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        "cache_read_input_tokens": 0,
+    }
+    assert not image.called
+
+
+@pytest.mark.parametrize("prompt", [{}, {"messages": []}, {"messages": "invalid"}])
+def test_sync_initial_estimate_fallback_does_not_consume_upstream(prompt: Mapping[str, object]) -> None:
+    wrapper: Final = AnthropicStreamWrapper(_FailingUpstream(), model="kimi-k2", completion_request=prompt)
+    first: Final = _EVENT.validate_python(next(wrapper))
+    message: Final = _EVENT.validate_python(first["message"])
+    assert message["usage"] == {
+        "input_tokens": 0,
+        "output_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        "cache_read_input_tokens": 0,
+    }
+    assert list(wrapper) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("prompt", [{}, {"messages": []}, {"messages": "invalid"}])
+async def test_async_initial_estimate_fallback_does_not_consume_upstream(prompt: Mapping[str, object]) -> None:
+    wrapper: Final = AnthropicStreamWrapper(_FailingUpstream(), model="kimi-k2", completion_request=prompt)
+    first: Final = _EVENT.validate_python(await anext(wrapper))
+    message: Final = _EVENT.validate_python(first["message"])
+    assert message["usage"] == {
+        "input_tokens": 0,
+        "output_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        "cache_read_input_tokens": 0,
+    }
+    with pytest.raises(RuntimeError, match="upstream failed before its first chunk"):
+        [chunk async for chunk in wrapper]
+
 
 MESSAGES = [{"role": "user", "content": "hello"}]
 

tests/unit/llms/bedrock/chat/test_invoke_handler.py

@@ -19,6 +19,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
 from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.llms.bedrock.chat.invoke_handler import (
+    AmazonAnthropicClaudeStreamDecoder,
     AmazonOpenAICompatibleStreamDecoder,
     AWSEventStreamDecoder,
     _stream_decoder,
@@ -1120,6 +1121,7 @@ async def test_converse_streaming_completion_yields_text_and_finish_reason(
             ),
         )
     )
+
     def handle_request(request: httpx.Request) -> httpx.Response:
         assert request.url.path.endswith("/converse-stream")
         return httpx.Response(200, content=event_stream_body, headers=_event_stream_headers())
@@ -1179,9 +1181,7 @@ def test_mistral_invoke_stream_yields_text_and_finish_reason(
     assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == (
         "A court case moves through the courts."
     )
-    assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == [
-        "stop"
-    ]
+    assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"]
 
 
 def test_converse_streaming_tool_call_arguments_are_valid_json(
@@ -1721,6 +1721,7 @@ def test_filter_headers_for_aws_signature():
 @pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
 def test_invoke_decoder_reports_only_missing_aws_dependency(missing):
     from unittest.mock import patch
+
     from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
 
     failure = ModuleNotFoundError("missing dependency", name=missing)
@@ -1835,9 +1836,7 @@ async def test_converse_stream_thinking_merges_into_content_with_think_tags() ->
             "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "c2lnbmF0dXJl"}}}
         ),
         _converse_event_frame("contentBlockStop", {"contentBlockIndex": 0}),
-        _converse_event_frame(
-            "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}}
-        ),
+        _converse_event_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}}),
         _converse_event_frame(
             "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "I am a helpful assistant"}}
         ),
@@ -1919,7 +1918,11 @@ _CUT_JSON_TOOL_STREAM: Final = b"".join(
             "content_block": {"type": "tool_use", "id": "toolu_cut_1", "name": "json_tool_call", "input": {}},
         },
         {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"city"'}},
-        {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ': "Paris", '}},
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "input_json_delta", "partial_json": ': "Paris", '},
+        },
         {
             "type": "content_block_delta",
             "index": 0,
@@ -1961,3 +1964,21 @@ async def test_async_anthropic_invoke_stream_cut_before_the_json_tool_block_ends
     chunks: Final = [chunk async for chunk in _cut_json_tool_stream_decoder().aiter_bytes(_chunks())]
 
     assert json.loads(_streamed_content(chunks)) == {"city": "Paris", "country": "France"}
+
+
+def test_invoke_stream_decoder_carries_the_guardrail_trace_of_the_message_stop_chunk() -> None:
+    """Bedrock Invoke reports the native guardrail trace once, on the final ``message_stop`` chunk, under
+    ``amazon-bedrock-trace``. The chunk keeps it so the stream builder and the spend log can price it."""
+    decoder = AmazonAnthropicClaudeStreamDecoder(model="anthropic.claude-sonnet-5-v1:0", sync_stream=True)
+    trace = {"guardrail": {"input": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 1}}}}}}
+    plain = decoder._chunk_parser({"type": "message_stop", "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3}})
+    traced = decoder._chunk_parser(
+        {
+            "type": "message_stop",
+            "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3},
+            "amazon-bedrock-trace": trace,
+        }
+    )
+    assert (plain.provider_specific_fields or {}).get("amazon-bedrock-trace") is None
+    assert traced.provider_specific_fields is not None
+    assert traced.provider_specific_fields["amazon-bedrock-trace"] == trace

tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py

@@ -4169,3 +4169,32 @@ async def test_bedrock_messages_400_is_not_retried_unless_it_rejects_opt_in_nati
 
     assert isinstance(raised, litellm.BadRequestError)
     assert len(bedrock.bodies) == 1
+
+
+@pytest.mark.parametrize("trace", ["enabled", "disabled"])
+def test_bedrock_messages_applies_the_deployment_guardrail_config_as_invoke_headers(trace: str) -> None:
+    from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import (
+        AmazonAnthropicClaudeMessagesConfig,
+    )
+
+    headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
+        headers={},
+        model="anthropic.claude-sonnet-5-v1:0",
+        messages=[],
+        optional_params={},
+        litellm_params={"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "2", "trace": trace}},
+    )
+    assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "gid"
+    assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "2"
+    assert headers["X-Amzn-Bedrock-Trace"] == trace.upper()
+
+
+def test_bedrock_messages_without_guardrail_config_sends_no_guardrail_headers() -> None:
+    from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import (
+        AmazonAnthropicClaudeMessagesConfig,
+    )
+
+    headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
+        headers={}, model="anthropic.claude-sonnet-5-v1:0", messages=[], optional_params={}, litellm_params={}
+    )
+    assert not any(name.startswith("X-Amzn-Bedrock-Guardrail") for name in headers)

tests/unit/llms/edenai/chat/test_edenai_chat_transformation.py

@@ -440,14 +440,3 @@ class TestDashboardRegistration:
         entry = json.loads(matrix_path.read_text())["providers"]["edenai"]
 
         assert entry["url"] == "https://docs.litellm.ai/docs/providers/edenai"
-        served = {name for name, flag in entry["endpoints"].items() if flag}
-        assert served == {
-            "chat_completions",
-            "messages",
-            "responses",
-            "embeddings",
-            "image_generations",
-            "audio_transcriptions",
-            "audio_speech",
-            "video_generations",
-        }

tests/unit/llms/openai_like/test_cortecs_provider.py

@@ -58,26 +58,6 @@ def test_cortecs_is_available_in_add_model_form():
     }
 
 
-def test_cortecs_supported_endpoints():
-    matrix_path = Path(litellm.__file__).parent / "provider_endpoints_support_backup.json"
-    providers = json.loads(matrix_path.read_text())["providers"]
-
-    assert providers["cortecs"]["endpoints"] == {
-        "chat_completions": True,
-        "messages": True,
-        "responses": True,
-        "embeddings": False,
-        "image_generations": False,
-        "audio_transcriptions": False,
-        "audio_speech": False,
-        "moderations": False,
-        "batches": False,
-        "rerank": False,
-        "a2a": False,
-        "interactions": False,
-    }
-
-
 def test_cortecs_chat_completion_request():
     with respx.mock() as upstream:
         route: Final = upstream.post("https://api.cortecs.ai/v1/chat/completions").respond(

tests/unit/llms/openai_like/test_libertai_provider.py

@@ -91,10 +91,3 @@ class TestLibertAIProviderConfig:
         matrix = json.loads(backup_path.read_text())
 
         assert "libertai" in matrix["providers"]
-        endpoints = matrix["providers"]["libertai"]["endpoints"]
-        assert endpoints["chat_completions"] is True
-        # embeddings is advertised false: the JSON-configured-provider path only
-        # wires chat routing (the OpenAILike embedding handler is reached solely
-        # for the literal openai_like/llamafile/lm_studio providers), matching
-        # the llamagate precedent. bge-m3 stays in the cost map for metadata.
-        assert endpoints["embeddings"] is False

tests/unit/llms/openai_like/test_prism_provider.py

@@ -99,25 +99,6 @@ def test_prism_is_available_in_add_model_form():
     }
 
 
-def test_prism_supported_endpoints():
-    matrix_path = Path(litellm.__file__).parent / "provider_endpoints_support_backup.json"
-    providers = json.loads(matrix_path.read_text())["providers"]
-
-    assert providers["prism"]["endpoints"] == {
-        "chat_completions": True,
-        "messages": True,
-        "responses": True,
-        "embeddings": False,
-        "image_generations": False,
-        "audio_transcriptions": False,
-        "audio_speech": False,
-        "moderations": False,
-        "batches": False,
-        "rerank": False,
-        "a2a": False,
-    }
-
-
 def test_prism_responses_request():
     with respx.mock() as upstream:
         route: Final = upstream.post("https://api.prisminference.com/v1/responses").respond(

tests/unit/llms/openai_like/test_reka_provider.py

@@ -115,28 +115,6 @@ def test_reka_is_available_in_add_model_form():
     }
 
 
-def test_reka_supported_endpoints():
-    expected: Final = {
-        "chat_completions": True,
-        "messages": True,
-        "responses": True,
-        "embeddings": False,
-        "image_generations": False,
-        "audio_transcriptions": False,
-        "audio_speech": False,
-        "moderations": False,
-        "batches": False,
-        "rerank": False,
-        "a2a": False,
-        "interactions": False,
-    }
-    backup_path = Path(litellm.__file__).parent / "provider_endpoints_support_backup.json"
-    root_path = Path(litellm.__file__).parent.parent / "provider_endpoints_support.json"
-
-    assert json.loads(backup_path.read_text())["providers"]["reka"]["endpoints"] == expected
-    assert json.loads(root_path.read_text())["providers"]["reka"]["endpoints"] == expected
-
-
 def test_reka_chat_completion_request():
     with respx.mock() as upstream:
         route: Final = upstream.post("https://api.reka.ai/v1/chat/completions").respond(200, json=_CHAT_COMPLETION)

tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py

@@ -1,7 +1,7 @@
 import asyncio
 import json
 from datetime import datetime
-from typing import Any, Dict, Final, List
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import httpx
@@ -284,9 +284,10 @@ class TestAzureAnthropicCostCalculation:
     @patch("litellm.completion_cost")
     def test_cost_calculation_with_azure_ai_custom_llm_provider(self, mock_completion_cost):
         """Test that custom_llm_provider is passed to completion_cost for Azure AI Anthropic"""
-        from litellm.types.utils import ModelResponse
         from datetime import datetime
 
+        from litellm.types.utils import ModelResponse
+
         mock_completion_cost.return_value = 0.001
 
         logging_obj = self._create_mock_logging_obj(
@@ -319,9 +320,10 @@ class TestAzureAnthropicCostCalculation:
     @patch("litellm.completion_cost")
     def test_cost_calculation_without_custom_llm_provider(self, mock_completion_cost):
         """Test that cost calculation works without custom_llm_provider (standard Anthropic)"""
-        from litellm.types.utils import ModelResponse
         from datetime import datetime
 
+        from litellm.types.utils import ModelResponse
+
         mock_completion_cost.return_value = 0.001
 
         # No custom_llm_provider in model_call_details
@@ -353,9 +355,10 @@ class TestAzureAnthropicCostCalculation:
     @patch("litellm.completion_cost")
     def test_cost_calculation_does_not_duplicate_provider_prefix(self, mock_completion_cost):
         """Test that provider prefix is not duplicated if already present in model name"""
-        from litellm.types.utils import ModelResponse
         from datetime import datetime
 
+        from litellm.types.utils import ModelResponse
+
         mock_completion_cost.return_value = 0.001
 
         logging_obj = self._create_mock_logging_obj(
@@ -826,9 +829,10 @@ class TestAnthropicBatchPassthroughCostTracking:
         mock_request_body,
     ):
         """Test that model gets 'anthropic/' prefix when not found in router"""
+        import base64
+
         from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig
         from litellm.types.utils import LiteLLMBatch
-        import base64
 
         # Model not in router - returns same model name
         mock_get_model_id.return_value = "claude-sonnet-4-5-20250929"
@@ -1089,7 +1093,7 @@ class TestAnthropicBatchPassthroughCostTracking:
 class TestBuildCompleteStreamingResponseRobustness:
     """_build_complete_streaming_response must tolerate non-standard SSE frames."""
 
-    def _build(self, chunks: List[str]):
+    def _build(self, chunks: list[str]):
         return AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
             all_chunks=chunks,
             litellm_logging_obj=MagicMock(),
@@ -1840,9 +1844,6 @@ class TestStreamFalseDeduplication:
         This is the _is_assembled_stream_success gate: with stream=False it
         always returned False and the guard was permanently disabled.
         """
-        from litellm.types.passthrough_endpoints.pass_through_endpoints import (
-            EndpointType,
-        )
         from litellm.types.utils import ModelResponse
 
         # Simulate what pass_through_endpoints.py now does after stream detection
@@ -2642,8 +2643,6 @@ def test_create_anthropic_response_logging_payload(mock_logging_obj, metadata_pa
 def test_handle_logging_anthropic_collected_chunks(all_chunks):
     from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
         AnthropicPassthroughLoggingHandler,
-        PassthroughStandardLoggingPayload,
-        EndpointType,
     )
     from litellm.types.utils import ModelResponse
 
@@ -2943,3 +2942,65 @@ def test_streamed_messages_from_other_providers_ignore_a_zero_usage_cost_field()
     headers = result["result"]._hidden_params.get("additional_headers") or {}
     assert "llm_provider-x-litellm-response-cost" not in headers
     assert result["kwargs"].get("response_cost") != 0.0
+
+
+_BEDROCK_TRACE_PRICING: Final = {
+    "bedrock/guardrails": {
+        "litellm_provider": "bedrock",
+        "mode": "guardrail",
+        "guardrail_cost_per_unit": {"topicPolicyUnits": 0.00015, "contentPolicyUnits": 0.00015},
+    }
+}
+_BEDROCK_INVOKE_TRACE: Final = {
+    "guardrail": {
+        "input": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}}}},
+        "outputs": [{"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 2, "contentPolicyUnits": 2}}}}],
+    }
+}
+
+
+def _recorded_guardrail_entries(logging_obj: LiteLLMLoggingObj) -> list:
+    return logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
+
+
+def test_anthropic_passthrough_handler_bills_the_bedrock_guardrail_trace_of_a_messages_response(
+    monkeypatch, mock_httpx_response, mock_response, mock_logging_obj
+) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING)
+    mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}}
+    AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler(
+        httpx_response=mock_httpx_response,
+        response_body={**mock_response, "amazon-bedrock-trace": _BEDROCK_INVOKE_TRACE},
+        logging_obj=mock_logging_obj,
+        url_route="/v1/messages",
+        result="success",
+        start_time=datetime(2025, 1, 1),
+        end_time=datetime(2025, 1, 1),
+        cache_hit=False,
+    )
+    (entry,) = _recorded_guardrail_entries(mock_logging_obj)
+    assert entry["guardrail_provider"] == "bedrock"
+    assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3}
+    assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015)
+
+
+def test_anthropic_passthrough_stream_bills_the_bedrock_guardrail_trace_of_the_message_stop_event(
+    monkeypatch, mock_logging_obj, all_chunks
+) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING)
+    mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}}
+    stop_event: Final = json.dumps({"type": "message_stop", "amazon-bedrock-trace": _BEDROCK_INVOKE_TRACE})
+    chunks: Final = [*all_chunks, f"event: message_stop\ndata: {stop_event}\n\n"]
+    AnthropicPassthroughLoggingHandler.handle_logging_anthropic_collected_chunks(
+        litellm_logging_obj=mock_logging_obj,
+        passthrough_success_handler_obj=MagicMock(),
+        url_route="/v1/messages",
+        request_body={"model": "claude", "messages": []},
+        endpoint_type=MagicMock(),
+        start_time=datetime(2025, 1, 1),
+        all_chunks=chunks,
+        end_time=datetime(2025, 1, 1),
+    )
+    (entry,) = _recorded_guardrail_entries(mock_logging_obj)
+    assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3}
+    assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015)

tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py

@@ -14,7 +14,7 @@ from openai.types.responses.response_function_web_search import (
     ActionSearchSource,
     ResponseFunctionWebSearch,
 )
-from pydantic import TypeAdapter
+from pydantic import TypeAdapter, ValidationError
 
 import litellm
 from litellm.constants import CALLER_SESSION_ID_METADATA_KEY
@@ -5062,6 +5062,50 @@ class TestHostedWebSearchReplay:
         assert replayed == search
         assert input_item == original
 
+    def test_replays_search_without_action(self) -> None:
+        input_item: Final = {"type": "web_search_call", "id": "ws_no_action", "status": "completed"}
+
+        messages: Final = (
+            LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
+                input_item=input_item
+            )
+        )
+
+        assert len(messages) == 1
+        assert messages[0]["role"] == "assistant"
+        content: Final = messages[0].get("content")
+        assert isinstance(content, str)
+        assert json.loads(content[content.index("{") :]) == input_item
+
+    @pytest.mark.parametrize(
+        "fields",
+        (
+            {"status": "completed"},
+            {"id": "ws_invalid"},
+            {"id": 42, "status": "completed"},
+            {"id": "ws_invalid", "status": "invalid"},
+            {"id": "ws_invalid", "status": "completed", "action": None},
+            {"id": "ws_invalid", "status": "completed", "action": 42},
+            {"id": "ws_invalid", "status": "completed", "action": {"type": "invalid"}},
+            {"id": "ws_invalid", "status": "completed", "action": {"type": "find_in_page"}},
+        ),
+        ids=(
+            "missing-id",
+            "missing-status",
+            "invalid-id",
+            "invalid-status",
+            "null-action",
+            "invalid-action-type",
+            "unknown-action",
+            "incomplete-action",
+        ),
+    )
+    def test_rejects_search_with_other_invalid_fields(self, fields: Mapping[str, object]) -> None:
+        with pytest.raises(ValidationError):
+            LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
+                input=[{"type": "web_search_call", **fields}], responses_api_request={}
+            )
+
     @pytest.mark.parametrize("order", ((0, 1, 2, 3), (1, 0, 3, 2), (1, 3, 0, 2)))
     @pytest.mark.parametrize("modify_params", (False, True))
     @pytest.mark.parametrize("structured_content", (False, True))