From 5667ac8497240541cd81d573ed1a7637456cd1f6 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 12:14:42 -0700 Subject: [PATCH 001/147] feat(proxy): refuse migrations that existing rows would make fail, before the first one runs (#48493) * feat(proxy): refuse migrations that existing rows would make fail, before the first one runs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(migrations): start the serving-proxy preflight case from a database without the api_key index Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type the preflight count rows without casts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format the db preflight flag Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy-extras): keep the extras files at their existing formatting Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy-extras): type the ledger query and the reader pid so the type gate holds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): declare the preflight report type where the extras package may be absent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): suppress the optional-extras unknown on the preflight report variables Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): suppress the optional-extras unknown on the preflight import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): return a typed preflight report when the extras package is unknown Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): sort the preflight import block Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark the preflight report binding Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(proxy): size the whole btree entry in the row check and name the flag for what it does The index check now lays out each row's entry the way index_form_tuple does (tuple header, null bitmap, every indexed column at its alignment, rounded to 8) and compares it with BTMaxItemSize for the server's block_size, so a key that passes octet_length > 2704 but still fails the (api_key, startTime) index is caught. --db_preflight becomes --check_migrations_against_existing_data, db_preflight.py becomes migration_data_checks.py and the timeout env var LITELLM_MIGRATION_DATA_CHECK_TIMEOUT Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration_data_checks.py | 639 ++++++++++++++++++ .../request_log_indexes.py | 12 + .../litellm_proxy_extras/utils.py | 226 +++---- litellm/proxy/db/prisma_client.py | 12 + litellm/proxy/proxy_cli.py | 27 + .../migrations/test_request_log_indexes.py | 201 +++++- .../test_migration_data_checks.py | 156 +++++ 7 files changed, 1144 insertions(+), 129 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py create mode 100644 tests/unit/litellm_proxy_extras/test_migration_data_checks.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py new file mode 100644 index 00000000000..a2d3728349a --- /dev/null +++ b/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 diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 6c31e8364a9..4c856c741d3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/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) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index be785e2fd78..528bf7def20 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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}" ) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 53971ed412c..3d783d8aba1 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/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 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 7f68b72a522..fc33ab1bd0b 100644 --- a/litellm/proxy/proxy_cli.py +++ b/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 " diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 3b6d1eaf21a..bf37ca80665 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/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) diff --git a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py new file mode 100644 index 00000000000..7897a388474 --- /dev/null +++ b/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 From 02605a14ada6c0e8cbc57258c28535e42c81a6c9 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 12:16:15 -0700 Subject: [PATCH 002/147] fix(anthropic): estimate initial input tokens for bridged message streams (#48618) * fix(anthropic): estimate initial input tokens for bridged message streams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): count iterable prompt blocks in initial usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): avoid fetching images during initial usage estimation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic/messages/adapters/handler.py | 2 + .../messages/adapters/streaming_iterator.py | 48 ++++-- .../messages/adapters/transformation.py | 2 + .../test_streaming_iterator_message_id.py | 150 ++++++++++++++++++ 4 files changed, 189 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index e5dd1d3c94f..71ccfbaf915 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/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: diff --git a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index 13acb19e454..d1853bb6bd5 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/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 diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 5ab112d0bfe..1ed552d46e3 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/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: diff --git a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py index b98877f25aa..ec9146bc023 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py +++ b/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"}] From 32297c917784a11309d08958b0bae8b86019235a Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 12:16:34 -0700 Subject: [PATCH 003/147] fix(logging): match extra_spend_tag_headers case-insensitively (#48620) * test(logging): cover extra spend tag header casing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): match spend tag headers case-insensitively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 12 ++-- .../test_litellm_logging.py | 67 +++++++++++++++++++ 2 files changed, 73 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 55f0f5d49f6..ee55c9ae030 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -7141,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 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 317d3c35664..a1f674b467f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/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 From 3d92124d9631bd6fef7edcb0ffdb08a445f242f9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:24:26 -0700 Subject: [PATCH 004/147] fix(vertex-ai): drop the floor retirement date from gemini-3.1-flash-lite-image (#48587) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 -- model_prices_and_context_window.json | 2 -- 2 files changed, 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b820b058eb3..390e9298489 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b820b058eb3..390e9298489 100644 --- a/model_prices_and_context_window.json +++ b/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, From 872bfd51024f9038af3cf196924da219a1e642a9 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 12:32:48 -0700 Subject: [PATCH 005/147] fix(proxy): return 403 naming the Enterprise licence for get_spend_routes without a licence (#48623) * test(proxy): cover enterprise spend permission errors on key and user creation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): return 403 naming the Enterprise licence for get_spend_routes without a licence Co-authored-by: er-s-an <92594860+er-s-an@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): narrow parsed permissions before checking spend access Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ci): allow full-history database checkout to finish --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: er-s-an <92594860+er-s-an@users.noreply.github.com> --- .github/workflows/test-unit.yml | 2 +- .../key_management_endpoints.py | 10 ++- .../test_key_management_endpoints.py | 79 +++++++++++++++++-- 3 files changed, 82 insertions(+), 9 deletions(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 9658e834c42..4d47a882891 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -952,7 +952,7 @@ jobs: dist: loadscope steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - timeout-minutes: 3 + timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }} with: persist-credentials: false fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }} diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b97d409a6c3..f7b72b263d4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -167,6 +167,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports handle_exception_on_proxy, hash_token_if_needed, is_valid_api_key, + require_enterprise_license, ) from litellm.repositories.base_repository import BaseRepository, record_to_dict from litellm.repositories.budget_repository import BudgetRepository @@ -5259,10 +5260,13 @@ async def generate_key_helper_fn( if isinstance(saved_token["metadata"], str): saved_token["metadata"] = json.loads(saved_token["metadata"]) if isinstance(saved_token["permissions"], str): - if "get_spend_routes" in saved_token["permissions"] and premium_user is not True: - raise ValueError("get_spend_routes permission is only available for LiteLLM Enterprise users") - saved_token["permissions"] = json.loads(saved_token["permissions"]) + if ( + isinstance(saved_token["permissions"], dict) + and "get_spend_routes" in saved_token["permissions"] + and premium_user is not True + ): + require_enterprise_license("get_spend_routes") if isinstance(saved_token["model_max_budget"], str): saved_token["model_max_budget"] = json.loads(saved_token["model_max_budget"]) router_settings = cast(dict | None, saved_token.get("router_settings")) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 4266c2ed6d1..9634f1efce8 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( + CommonProxyErrors, GenerateKeyRequest, GenerateKeyResponse, KeyManagementRoutes, @@ -21354,8 +21355,9 @@ async def test_key_generation_premium_permission_error_precedes_project_ownershi team_table=None, ) - assert error.value.status_code == 500 - assert error.value.detail == {"error": "Internal Server Error."} + assert error.value.status_code == 403 + assert "get_spend_routes" in str(error.value.detail) + assert CommonProxyErrors.not_premium_user.value in str(error.value.detail).replace("\\n", "\n") mock_prisma_client.insert_data.assert_not_awaited() @@ -23210,6 +23212,76 @@ def test_key_info_unexpected_database_exception_sanitized( proxy_server.app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.fixture +def spend_permission_client(monkeypatch: pytest.MonkeyPatch) -> TestClient: + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + prisma_client: Final = _configure_key_endpoints(monkeypatch, UserApiKeyCache()) + prisma_client.db.litellm_usertable = MagicMock() + prisma_client.db.litellm_usertable.count = AsyncMock(return_value=0) + monkeypatch.setattr(litellm, "default_key_generate_params", None) + monkeypatch.setattr(proxy_server, "premium_user", False) + monkeypatch.setitem( + proxy_server.app.dependency_overrides, + user_api_key_auth, + lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + ) + return TestClient(proxy_server.app) + + +@pytest.mark.parametrize("route", ("/key/generate", "/user/new")) +@pytest.mark.parametrize("permission", ("true", "false", False), ids=("string_true", "string_false", "bool_false")) +def test_spend_permission_enterprise_error_regression( + spend_permission_client: TestClient, + route: str, + permission: str | bool, +) -> None: + response: Final = spend_permission_client.post(route, json={"permissions": {"get_spend_routes": permission}}) + + assert response.status_code == 403, response.text + assert "get_spend_routes" in response.text + assert CommonProxyErrors.not_premium_user.value in response.json()["error"]["message"].replace("\\n", "\n") + + +@pytest.mark.parametrize( + "permissions", + ({}, {"other": True}, {"other": "get_spend_routes"}, {"get_spend_routes_extra": True}), + ids=("empty", "unrelated", "permission_in_value", "permission_in_other_key"), +) +def test_non_enterprise_permissions_guard( + spend_permission_client: TestClient, + permissions: dict[str, str | bool], +) -> None: + response: Final = spend_permission_client.post("/key/generate", json={"permissions": permissions}) + + assert response.status_code == 200, response.text + assert response.json()["permissions"] == permissions + assert response.json()["key"].startswith("sk-") + + +def test_enterprise_spend_permission_guard( + spend_permission_client: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + response: Final = spend_permission_client.post( + "/key/generate", json={"permissions": {"get_spend_routes": True}} + ) + + assert response.status_code == 200, response.text + assert response.json()["permissions"] == {"get_spend_routes": True} + assert response.json()["key"].startswith("sk-") + + +def test_guardrails_enterprise_error_guard(spend_permission_client: TestClient) -> None: + response: Final = spend_permission_client.post("/key/generate", json={"guardrails": ["g"]}) + + assert response.status_code == 403, response.text + assert "guardrails" in response.text + assert CommonProxyErrors.not_premium_user.value in response.json()["error"]["message"].replace("\\n", "\n") + + def test_key_info_missing_key_preserves_404( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -23300,6 +23372,3 @@ def test_key_info_database_not_connected_sanitized_500( assert "Database not connected" in caplog.text finally: proxy_server.app.dependency_overrides.pop(user_api_key_auth, None) - - - From 2dc4847579bb9ab4b8bff45b49022b2752658f3b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 12:33:48 -0700 Subject: [PATCH 006/147] fix(token_counter): count Anthropic-format tool_choice without KeyError (#48624) * test(token_counter): cover Anthropic tool_choice counting Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(token_counter): count Anthropic-format tool_choice without KeyError Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/token_counter.py | 23 ++++++++----- .../litellm_core_utils/test_token_counter.py | 34 ++++++++++++++++++- .../proxy/spend_tracking/test_input_tokens.py | 26 ++++++++++++++ 3 files changed, 73 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index ee5f70fb394..edb5ecf0f8d 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -59,7 +59,6 @@ from litellm.types.llms.anthropic import ( from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionDocumentObject, - ChatCompletionNamedToolChoiceParam, ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock, ChatCompletionToolParam, @@ -456,7 +455,7 @@ def token_counter( messages: Sequence[AllMessageValues | Message] | None = None, count_response_tokens: bool | None = False, tools: list[ChatCompletionToolParam] | None = None, - tool_choice: ChatCompletionNamedToolChoiceParam | None = None, + tool_choice: str | Mapping[str, object] | None = None, use_default_image_token_count: bool | None = False, default_token_count: int | None = None, ) -> int: @@ -636,7 +635,7 @@ def _count_messages( def _count_extra( count_function: TokenCounterFunction, tools: list[ChatCompletionToolParam] | None, - tool_choice: ChatCompletionNamedToolChoiceParam | None, + tool_choice: str | Mapping[str, object] | None, includes_system_message: bool, tool_section: ChatTemplateToolSection | None, ) -> int: @@ -673,13 +672,19 @@ def _count_tool_definitions( def _count_tool_choice( count_function: TokenCounterFunction, - tool_choice: ChatCompletionNamedToolChoiceParam | None, + tool_choice: str | Mapping[str, object] | None, ) -> int: - if tool_choice == "none": - return 1 - if isinstance(tool_choice, dict): - return 7 + count_function(str(tool_choice["function"]["name"])) - return 0 + match tool_choice: + case "none": + return 1 + case {"function": {"name": str() as name}}: + return 7 + count_function(name) + case {"type": "none"}: + return 1 + case {"name": str() as name}: + return 7 + count_function(name) + case _: + return 0 def _count_chat_template_tools( diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index f222a257b8f..9921bbe343a 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import Future, ThreadPoolExecutor, wait from itertools import accumulate, chain from pathlib import Path -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, sentinel import anyio.to_thread @@ -42,6 +42,7 @@ from litellm.litellm_core_utils.token_counter import ( ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from tests.large_text import text from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, @@ -658,6 +659,37 @@ def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: obj ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) +@pytest.mark.parametrize( + ("tool_choice", "expected_tool_choice"), + [ + pytest.param({"type": "auto"}, None, id="anthropic-auto"), + pytest.param({"type": "any"}, None, id="anthropic-any"), + pytest.param({"type": "none"}, "none", id="anthropic-none"), + pytest.param( + {"type": "tool", "name": "get_weather"}, + {"type": "function", "function": {"name": "get_weather"}}, + id="anthropic-tool", + ), + pytest.param( + {"type": "function", "name": "get_weather"}, + {"type": "function", "function": {"name": "get_weather"}}, + id="responses-function", + ), + ], +) +def test_token_counter_counts_anthropic_tool_choices_like_openai_choices( + tool_choice: Mapping[str, object], expected_tool_choice: str | Mapping[str, object] | None +) -> None: + messages: Final = cast(list[AllMessageValues], list(_USER_MESSAGES)) + tools: Final = cast(list[ChatCompletionToolParam], [_openai_tool(_WEATHER_FUNCTION)]) + + assert token_counter_new( + model="gpt-4", messages=messages, tools=tools, tool_choice=tool_choice + ) == token_counter_new( + model="gpt-4", messages=messages, tools=tools, tool_choice=expected_tool_choice + ) + + @pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) @pytest.mark.parametrize( "anthropic_tool", diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index 71f4dec2345..c7f13651471 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -30,6 +30,32 @@ async def test_large_input_is_still_counted() -> None: assert isinstance(counts, MappingProxyType) +def test_count_input_tokens_for_model_counts_anthropic_auto_tool_choice() -> None: + model: Final = "claude-sonnet-5-5" + body_without_tool_choice: Final = { + "model": model, + "messages": [{"role": "user", "content": "weather in pune?"}], + "tools": [ + { + "name": "get_weather", + "description": "Get the weather for a location", + "input_schema": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + } + ], + } + body_with_tool_choice: Final = body_without_tool_choice | {"tool_choice": {"type": "auto"}} + + expected: Final = count_input_tokens_for_model(request_body=body_without_tool_choice, model=model) + actual: Final = count_input_tokens_for_model(request_body=body_with_tool_choice, model=model) + + assert isinstance(expected, int) + assert actual == expected + + _COMPACTION_ITEM: Final = { "type": "compaction", "id": "cmp_1", From 189ab854389f49017c7568571de19b964b8e8e01 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:38:14 +0000 Subject: [PATCH 007/147] fix(bedrock): correct gpt max_input_tokens and add Mantle token counting for claude-opus-5-5 (#48479) * fix(bedrock): set gpt-6.1-terra max_input_tokens to the 1,050,000 Bedrock card context window * fix(bedrock): set gpt-5.x and gpt-6 max_input_tokens to the card context window minus max output * fix(bedrock): add Mantle token_count_model to bedrock_mantle claude-opus-5-5 rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: kerry --- ...odel_prices_and_context_window_backup.json | 107 ++++++++++++------ model_prices_and_context_window.json | 107 ++++++++++++------ .../common_utils/test_token_count_settings.py | 2 +- 3 files changed, 149 insertions(+), 67 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 390e9298489..7cd4eadca6e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57957,7 +57957,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58000,7 +58000,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58108,7 +58108,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58148,7 +58148,7 @@ "output_cost_per_token_above_272k_tokens": 3.3e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-sol.html", @@ -58184,7 +58184,7 @@ "output_cost_per_token_above_272k_tokens": 3e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-sol.html", @@ -58220,7 +58220,7 @@ "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", @@ -58256,7 +58256,7 @@ "output_cost_per_token_above_272k_tokens": 1.8e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", @@ -58292,7 +58292,7 @@ "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", @@ -58325,7 +58325,7 @@ "output_cost_per_token_above_272k_tokens": 2.475e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", @@ -58356,7 +58356,7 @@ "output_cost_per_token_above_272k_tokens": 2.25e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", @@ -58387,7 +58387,7 @@ "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", @@ -58418,7 +58418,7 @@ "output_cost_per_token_above_272k_tokens": 4.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", @@ -58452,7 +58452,7 @@ "output_cost_per_token_above_272k_tokens": 1.8e-06, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", @@ -58492,7 +58492,7 @@ "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, @@ -58532,7 +58532,7 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58570,7 +58570,7 @@ "output_cost_per_token": 5.5e-07, "output_cost_per_token_above_272k_tokens": 8.25e-07, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58616,7 +58616,7 @@ "output_cost_per_token_above_272k_tokens": 8.25e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, @@ -58656,7 +58656,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58694,7 +58694,7 @@ "output_cost_per_token_above_272k_tokens": 8.25e-07, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58738,7 +58738,7 @@ "output_cost_per_token_above_272k_tokens": 7.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token_above_272k_tokens_ultrafast": 0.00045, @@ -58777,7 +58777,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58811,7 +58811,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58848,7 +58848,7 @@ "output_cost_per_token_above_272k_tokens": 7.5e-07, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58882,7 +58882,7 @@ "output_cost_per_token_above_272k_tokens": 7.5e-07, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58921,7 +58921,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58960,7 +58960,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -59181,7 +59181,8 @@ "prompt_cache_min_tokens": 512, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-5-5.html", "thinking_always_on": true, - "supports_forced_tool_use": false + "supports_forced_tool_use": false, + "token_count_model": {"model":"anthropic.claude-opus-5-5","litellm_provider":"bedrock_mantle","region":"us-east-1"} }, "bedrock_mantle/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, @@ -66036,7 +66037,7 @@ }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66071,7 +66072,7 @@ }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66106,7 +66107,7 @@ }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66297,7 +66298,8 @@ "supports_xhigh_reasoning_effort": true, "thinking_always_on": true, "supports_forced_tool_use": false, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-5-5.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-5-5.html", + "token_count_model": {"model":"anthropic.claude-opus-5-5","litellm_provider":"bedrock_mantle","region":"us-gov-west-1"} }, "bedrock_mantle/us-gov-west-1/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, @@ -66336,7 +66338,7 @@ }, "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66360,7 +66362,8 @@ "supports_vision": true, "input_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 3.3e-07, - "output_cost_per_token": 1.98e-05 + "output_cost_per_token": 1.98e-05, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.md" }, "bedrock_mantle/us-gov-east-1/xai.grok-4.6": { "use_openai_responses_path": true, @@ -76003,5 +76006,43 @@ "high", "max" ] + }, + "bedrock_mantle/openai.gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "use_openai_responses_path": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 390e9298489..7cd4eadca6e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57957,7 +57957,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58000,7 +58000,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58108,7 +58108,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58148,7 +58148,7 @@ "output_cost_per_token_above_272k_tokens": 3.3e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-sol.html", @@ -58184,7 +58184,7 @@ "output_cost_per_token_above_272k_tokens": 3e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-sol.html", @@ -58220,7 +58220,7 @@ "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", @@ -58256,7 +58256,7 @@ "output_cost_per_token_above_272k_tokens": 1.8e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", @@ -58292,7 +58292,7 @@ "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", @@ -58325,7 +58325,7 @@ "output_cost_per_token_above_272k_tokens": 2.475e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", @@ -58356,7 +58356,7 @@ "output_cost_per_token_above_272k_tokens": 2.25e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", @@ -58387,7 +58387,7 @@ "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", @@ -58418,7 +58418,7 @@ "output_cost_per_token_above_272k_tokens": 4.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", @@ -58452,7 +58452,7 @@ "output_cost_per_token_above_272k_tokens": 1.8e-06, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", @@ -58492,7 +58492,7 @@ "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, @@ -58532,7 +58532,7 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58570,7 +58570,7 @@ "output_cost_per_token": 5.5e-07, "output_cost_per_token_above_272k_tokens": 8.25e-07, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58616,7 +58616,7 @@ "output_cost_per_token_above_272k_tokens": 8.25e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, @@ -58656,7 +58656,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58694,7 +58694,7 @@ "output_cost_per_token_above_272k_tokens": 8.25e-07, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58738,7 +58738,7 @@ "output_cost_per_token_above_272k_tokens": 7.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token_above_272k_tokens_ultrafast": 0.00045, @@ -58777,7 +58777,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58811,7 +58811,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58848,7 +58848,7 @@ "output_cost_per_token_above_272k_tokens": 7.5e-07, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58882,7 +58882,7 @@ "output_cost_per_token_above_272k_tokens": 7.5e-07, "litellm_provider": "bedrock_converse", "supports_bedrock_converse_tool_result_images": false, - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "supported_modalities": [ @@ -58921,7 +58921,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -58960,7 +58960,7 @@ "search_context_size_medium": 0.012 }, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -59181,7 +59181,8 @@ "prompt_cache_min_tokens": 512, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-5-5.html", "thinking_always_on": true, - "supports_forced_tool_use": false + "supports_forced_tool_use": false, + "token_count_model": {"model":"anthropic.claude-opus-5-5","litellm_provider":"bedrock_mantle","region":"us-east-1"} }, "bedrock_mantle/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, @@ -66036,7 +66037,7 @@ }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66071,7 +66072,7 @@ }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66106,7 +66107,7 @@ }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66297,7 +66298,8 @@ "supports_xhigh_reasoning_effort": true, "thinking_always_on": true, "supports_forced_tool_use": false, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-5-5.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-5-5.html", + "token_count_model": {"model":"anthropic.claude-opus-5-5","litellm_provider":"bedrock_mantle","region":"us-gov-west-1"} }, "bedrock_mantle/us-gov-west-1/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, @@ -66336,7 +66338,7 @@ }, "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, @@ -66360,7 +66362,8 @@ "supports_vision": true, "input_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 3.3e-07, - "output_cost_per_token": 1.98e-05 + "output_cost_per_token": 1.98e-05, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.md" }, "bedrock_mantle/us-gov-east-1/xai.grok-4.6": { "use_openai_responses_path": true, @@ -76003,5 +76006,43 @@ "high", "max" ] + }, + "bedrock_mantle/openai.gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "use_openai_responses_path": true } } diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index d7314f1a6d3..366d8b010d0 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -183,7 +183,7 @@ def test_every_shipped_token_count_model_resolves_to_a_mapped_aws_model() -> Non for name, count_model in count_models.items(): model: Final = count_model["model"] assert get_bedrock_base_model(model) in aws_models - assert cost_map[name]["litellm_provider"] in ("bedrock", "bedrock_converse") + assert cost_map[name]["litellm_provider"] in ("bedrock", "bedrock_converse", "bedrock_mantle") assert "region" not in count_model or count_model["region"].strip() From 788868a52f1901d8c1570e58fb79014d572647c2 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 27 Oct 2026 12:40:59 -0700 Subject: [PATCH 008/147] fix(security)!: build proxy-only failure logs from an allowlist and drop caller guardrail info (#48599) * fix(proxy): build proxy-only failure logs from an allowlist and ignore request-time GCS credentials * fix(proxy): drop caller metadata from proxy-only failure logs metadata and litellm_metadata leave the proxy-only failure-log allowlist, so a rejected request can no longer point its failure log at its own Langsmith or Langfuse host through nested metadata. * fix(proxy): keep allowlisted metadata in proxy-only failure logs * test(proxy): isolate proxy-only failure log metadata tests * fix(logging): ignore non-list request tags when building the logging payload A guardrail-rejected request with metadata tags set to a non-list value crashed the standard logging payload, so its spend row lost model attribution. * fix(proxy): keep router attribution and drop unauthenticated tags in proxy-only failure logs A request rejected before routing now logs the configured deployment and model group even when its metadata or litellm_metadata carries its own values. Requests rejected at auth no longer write caller tags into failure logs and tag spend. * fix(proxy): always strip client guardrail info and mirror failure attribution into litellm_metadata allow_client_pricing_override no longer keeps a caller-supplied standard_logging_guardrail_information on the chat or pass-through paths, since its guardrail_cost is billed. Proxy-only failure logs now also carry key and team attribution in litellm_metadata. * refactor(proxy): keep the proxy-only failure-log fixes inside the type gates Types the picked helpers so basedpyright and the type-discipline gate stay at main's counts. No runtime change. * fix(proxy): keep the WS frame guardrail-info strip and string interaction tags as sent The guardrail-info key left the pricing strip list, so the Responses WS frame guard now drops it by name. get_request_tags no longer raises on string tags, so background interaction polling checks for them before keeping them as sent. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../gcs_bucket/gcs_bucket_base.py | 49 +----- .../interactions/background_cost_polling.py | 32 +++- litellm/litellm_core_utils/litellm_logging.py | 20 +-- litellm/proxy/litellm_pre_call_utils.py | 19 ++- .../pass_through_endpoints.py | 8 +- .../proxy/response_api_endpoints/endpoints.py | 2 +- litellm/proxy/utils.py | 114 +++++++++++-- .../gcs_bucket/test_gcs_bucket_base.py | 18 ++ .../test_litellm_logging.py | 15 ++ .../files/test_vertex_ai_files_handler.py | 84 ---------- .../test_pass_through_endpoints.py | 34 ++++ tests/unit/proxy/test_pricing_field_strip.py | 31 +++- ..._utils_model_creation_and_error_logging.py | 157 +++++++++++++----- .../test_post_call_failure_hook.py | 122 +++++++++++++- 14 files changed, 494 insertions(+), 211 deletions(-) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_base.py b/litellm/integrations/gcs_bucket/gcs_bucket_base.py index 63c4fe92eca..092372b5d15 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_base.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_base.py @@ -20,7 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.integrations.gcs_bucket import * -from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload +from litellm.types.utils import StandardLoggingPayload if TYPE_CHECKING: from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -140,49 +140,16 @@ class GCSBucketBase(CustomBatchLogger): async def get_gcs_logging_config(self, kwargs: dict[str, Any] | None = {}) -> GCSLoggingConfig: """ - This function is used to get the GCS logging config for the GCS Bucket Logger. - It checks if the dynamic parameters are provided in the kwargs and uses them to get the GCS logging config. - If no dynamic parameters are provided, it uses the default values. + Returns the bucket and service account this logger was configured with. """ - if kwargs is None: - kwargs = {} - - standard_callback_dynamic_params: Final[StandardCallbackDynamicParams | None] = kwargs.get( - "standard_callback_dynamic_params", None - ) - - bucket_name: str - path_service_account: str | None - if standard_callback_dynamic_params is not None: - verbose_logger.debug("Using dynamic GCS logging") - verbose_logger.debug("standard_callback_dynamic_params: %s", standard_callback_dynamic_params) - - _bucket_name: str | None = standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME - _path_service_account: Final[str | None] = ( - standard_callback_dynamic_params.get("gcs_path_service_account", None) or self.path_service_account_json + if self.BUCKET_NAME is None: + raise ValueError( + "GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment." ) - - if _bucket_name is None: - raise ValueError( - "GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment." - ) - bucket_name = _bucket_name - path_service_account = _path_service_account - vertex_instance = await self.get_or_create_vertex_instance(credentials=path_service_account) - else: - # If no dynamic parameters, use the default instance - if self.BUCKET_NAME is None: - raise ValueError( - "GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment." - ) - bucket_name = self.BUCKET_NAME - path_service_account = self.path_service_account_json - vertex_instance = await self.get_or_create_vertex_instance(credentials=path_service_account) - return GCSLoggingConfig( - bucket_name=bucket_name, - vertex_instance=vertex_instance, - path_service_account=path_service_account, + bucket_name=self.BUCKET_NAME, + vertex_instance=await self.get_or_create_vertex_instance(credentials=self.path_service_account_json), + path_service_account=self.path_service_account_json, ) async def get_or_create_vertex_instance(self, credentials: str | None) -> VertexBase: diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 275eb17787e..0b8c128286a 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -313,19 +313,35 @@ def _spend_metadata(logging_obj: "LiteLLMLoggingObj") -> Mapping[str, object]: } +def _sent_tags_are_a_list(litellm_params: Mapping[str, object]) -> bool: + """Whether the first non-empty ``tags`` in ``metadata`` then ``litellm_metadata`` is absent or a list.""" + buckets: Final = (litellm_params.get("metadata"), litellm_params.get("litellm_metadata")) + sent: Final = next( + ( + tags + for bucket in buckets + if isinstance(bucket, Mapping) and (tags := _OBJECT_MAPPING.validate_python(bucket).get("tags")) + ), + None, + ) + return sent is None or isinstance(sent, list) + + def _request_tags(litellm_params: Mapping[str, object]) -> Mapping[str, object]: from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup try: - return { - "tags": StandardLoggingPayloadSetup.get_request_tags( - litellm_params=_OBJECT_DICT.validate_python(litellm_params), - proxy_server_request=_OBJECT_DICT.validate_python(litellm_params.get("proxy_server_request") or {}), - ) - } + if _sent_tags_are_a_list(litellm_params): + return { + "tags": StandardLoggingPayloadSetup.get_request_tags( + litellm_params=_OBJECT_DICT.validate_python(litellm_params), + proxy_server_request=_OBJECT_DICT.validate_python(litellm_params.get("proxy_server_request") or {}), + ) + } except AttributeError: - verbose_logger.debug("Keeping a background interaction's request tags as sent: they are not a list") - return {} + pass + verbose_logger.debug("Keeping a background interaction's request tags as sent: they are not a list") + return {} def _create_context(logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str) -> BackgroundInteractionCreateContext: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ee55c9ae030..48e8a44b9a4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -7161,19 +7161,17 @@ class StandardLoggingPayloadSetup: litellm_metadata: Final = cast( # cast-ok: request metadata is caller-provided Mapping[str, object], litellm_params.get("litellm_metadata") or {} ) - if metadata.get("tags", []): - request_tags = cast(list[str], metadata.get("tags", [])).copy() # cast-ok: tags are caller-provided - elif litellm_metadata.get("tags", []): - request_tags = cast(list[str], litellm_metadata.get("tags", [])).copy() # cast-ok: tags are caller-provided - else: - request_tags = [] + request_tags: Final = next( + ( + cast(list[str], tags) # cast-ok: tags are caller-provided + for tags in (metadata.get("tags"), litellm_metadata.get("tags")) + if isinstance(tags, list) and tags + ), + (), + ) user_agent_tags: Final = StandardLoggingPayloadSetup._get_user_agent_tags(proxy_server_request) additional_header_tags: Final = StandardLoggingPayloadSetup._get_extra_header_tags(proxy_server_request) - if user_agent_tags is not None: - request_tags.extend(user_agent_tags) - if additional_header_tags is not None: - request_tags.extend(additional_header_tags) - return request_tags + return [*request_tags, *(user_agent_tags or ()), *(additional_header_tags or ())] _get_request_tags = get_request_tags diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index baa150efc86..1e7cdc44c70 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -466,10 +466,7 @@ _ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY: Final = "allow_client_mess _CLIENT_PRICING_CONTROL_FIELDS: Final = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) | {"base_model"} # ``model_info`` carries the same pricing fields when read by # ``use_custom_pricing_for_model``; strip from metadata for the same reason. -# ``standard_logging_guardrail_information`` is proxy-written telemetry summed -# into response_cost and spend; a client seeding it forges (even negative) -# guardrail cost. -_CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logging_guardrail_information"}) +_CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info"}) # ``attempted_fallbacks`` and ``original_model_group`` are written by the router # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. @@ -799,6 +796,19 @@ def _strip_client_pricing_overrides(data: dict[str, object]) -> None: ) +def strip_client_guardrail_information(data: dict[str, object]) -> None: # mutable-ok: strips the request body in place + """Guardrail hooks write this field and its cost is billed, so no key or team setting keeps a caller's value.""" + for metadata_key in ("metadata", "litellm_metadata"): + if ( + isinstance(metadata := data.get(metadata_key), dict) + and "standard_logging_guardrail_information" in metadata + ): + del metadata["standard_logging_guardrail_information"] + verbose_proxy_logger.debug( + "Stripped client-supplied guardrail information from request body: %s", metadata_key + ) + + def _strip_router_reserved_metadata( data: dict[str, object], # mutable-ok: strips in place on the request body the pre-call pipeline threads through ) -> None: @@ -2695,6 +2705,7 @@ async def add_litellm_data_to_request( # would silently skip the field. if not _key_or_team_allows_client_pricing_override(user_api_key_dict): _strip_client_pricing_overrides(data) + strip_client_guardrail_information(data) _strip_router_reserved_metadata(data) # Same reason as the strips above: runs after the metadata string-to-dict parse diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fc0e5be3f4c..97f2c10afb1 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -119,6 +119,7 @@ from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy modul _strip_client_pricing_overrides, # pyright: ignore[reportPrivateUsage] # sanitize before trusted hooks add guardrail costs get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above redact_credential_headers, + strip_client_guardrail_information, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.route_priority import ( @@ -1223,10 +1224,11 @@ async def pass_through_request( _parsed_body, ) + client_body: Final = TypeAdapter(dict[str, object]).validate_python(_parsed_body) + strip_client_guardrail_information(client_body) if not _key_or_team_allows_client_pricing_override(user_api_key_dict): - pricing_body: Final = TypeAdapter(dict[str, object]).validate_python(_parsed_body) - _strip_client_pricing_overrides(pricing_body) - _parsed_body = pricing_body + _strip_client_pricing_overrides(client_body) + _parsed_body = client_body if custom_llm_provider in ("laya", "bespoke"): decision_request: Final = TypeAdapter(Mapping[str, object]).validate_python(_parsed_body) checkpoint: Final = validate_oss_request(custom_llm_provider, decision_request) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 312c5124ca5..97bbf893e3e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2116,7 +2116,7 @@ def _without_untrusted_metadata(key: str, value: object) -> object: if not name.startswith("user_api_key_") and name not in _UNTRUSTED_METADATA_CONTROL_FIELDS and name not in _CLIENT_PRICING_METADATA_FIELDS - and name not in ("headers", "pii_tokens", *_CLIENT_GUARDRAIL_SELECTORS) + and name not in ("headers", "pii_tokens", _GUARDRAIL_INFORMATION_KEY, *_CLIENT_GUARDRAIL_SELECTORS) } diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ce7e6afde24..d3202368d76 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1387,7 +1387,8 @@ def _stamp_deployment_attribution( callbacks key off it and must not count a proxy-side reject against the deployment. A failure after the provider handoff keeps the metadata the router stamped; a request that never reached a provider is flagged ``PROXY_REJECTED_BEFORE_ROUTING_KEY`` (deployment metrics key off it) whatever - its metadata says, since ``metadata.model_info`` can be caller supplied.""" + its metadata says, since ``metadata.model_info`` can be caller supplied. Its ``deployment`` and + ``model_group`` always come from the router config, never from the caller's metadata.""" attribution: Final = _deployment_attribution_for_model_group(model_group, team_id) if "custom_llm_provider" in attribution: litellm_params["custom_llm_provider"] = attribution["custom_llm_provider"] @@ -1402,9 +1403,13 @@ def _stamp_deployment_attribution( if not isinstance(metadata, dict): return attribution metadata.setdefault("model_info", attribution["model_info"]) - metadata.setdefault("deployment", attribution["deployment"]) - if isinstance(model_group, str): - metadata.setdefault("model_group", model_group) + routing: Final = MappingProxyType( + {"deployment": attribution["deployment"], **({"model_group": model_group} if model_group else {})} + ) + metadata.update(routing) + litellm_metadata: Final = litellm_params.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + litellm_metadata.update(routing) return attribution @@ -1484,6 +1489,9 @@ object rather than ``True`` because ``request_data`` is the client body: JSON ca client can neither hide a guardrail block nor force a row by sending the key or a ``litellm_call_id`` itself.""" _PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) +_PROXY_ONLY_ERROR_LOGGED_FIELDS: Final = frozenset( + {"model", "messages", "input", "prompt", "user", "stream", "litellm_call_id"} +) _LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data")) @@ -1518,6 +1526,57 @@ def _provider_billed_the_failure(exception: Exception) -> bool: return _timed_out_reading_from_provider(exception) or _client_disconnected(exception) +def _as_str_mapping(value: object) -> Mapping[str, object]: + if not isinstance(value, Mapping): + return MappingProxyType({}) + return cast( # cast-ok: the runtime Mapping check validates the container shape + Mapping[str, object], value + ) + + +def _proxy_only_error_logged_metadata_keys(*, authenticated: bool, past_pre_call: bool) -> frozenset[str]: + return frozenset( + key + for key, logged in (("tags", authenticated), ("standard_logging_guardrail_information", past_pre_call)) + if logged + ) + + +def _is_logged_metadata_value(key: str, value: object) -> bool: + if key != "tags": + return True + tags: Final = cast(list[object], value) if isinstance(value, list) else None # cast-ok: isinstance checked the list + return tags is not None and all(isinstance(item, str) for item in tags) + + +def _proxy_only_error_log_metadata( + value: object, logged_keys: frozenset[str] +) -> dict[str, object] | None: # mutable-ok: router attribution stamps the logged metadata in place + metadata: Final = _as_str_mapping(value) + if not isinstance(value, Mapping): + return None + return { + key: value for key, value in metadata.items() if key in logged_keys and _is_logged_metadata_value(key, value) + } + + +def _proxy_only_error_log_fields( + request_data: Mapping[str, object], logged_metadata_keys: frozenset[str] +) -> Mapping[str, object]: + metadata_fields: Final = { + key: filtered_metadata + for key in ("metadata", "litellm_metadata") + if (filtered_metadata := _proxy_only_error_log_metadata(request_data.get(key), logged_metadata_keys)) + is not None + } + return MappingProxyType( + { + **{key: value for key, value in request_data.items() if key in _PROXY_ONLY_ERROR_LOGGED_FIELDS}, + **metadata_fields, + } + ) + + def _failure_fields_to_lift( request_data: Mapping[str, object], provider_billed: bool = False, dispatched_call_type: str | None = None ) -> Mapping[str, object]: @@ -3778,6 +3837,7 @@ class ProxyLogging: user_api_key_dict=user_api_key_dict, route=route, original_exception=original_exception, + error_type=error_type, ) request_data.update( @@ -3872,6 +3932,7 @@ class ProxyLogging: user_api_key_dict: UserAPIKeyAuth, route: str | None = None, original_exception: Exception | None = None, + error_type: ProxyErrorTypes | None = None, ): """ Handle logging for proxy only errors by calling `litellm_logging_obj.async_failure_handler` @@ -3880,19 +3941,25 @@ class ProxyLogging: """ litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None) set_up_by_request: Final = litellm_logging_obj is not None + request_fields: Final = cast(Mapping[str, object], request_data) # cast-ok: request bodies are str-keyed JSON + logged_metadata_keys: Final = _proxy_only_error_logged_metadata_keys( + authenticated=error_type != ProxyErrorTypes.auth_error, past_pre_call=set_up_by_request + ) + user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict + ) if litellm_logging_obj is None: from litellm._uuid import uuid request_data.setdefault("litellm_call_id", str(uuid.uuid4())) - user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) litellm_logging_obj, data = litellm.utils.function_setup( original_function=route or "IGNORE_THIS", rules_obj=litellm.utils.Rules(), start_time=datetime.now(), - **request_data, + **_proxy_only_error_log_fields( # pyright: ignore[reportArgumentType] # the allowlist never carries is_async_call + request_fields, logged_metadata_keys + ), ) request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP}) @@ -3902,18 +3969,39 @@ class ProxyLogging: if litellm_logging_obj is not None: ## UPDATE LOGGING INPUT - _optional_params: Final = {} - _litellm_params: Final = {} + _optional_params: Final[dict[str, object]] = {} # mutable-ok: filled below, then handed to the logger + _litellm_params: Final[dict[str, object]] = {} # mutable-ok: filled below, then handed to the logger litellm_param_keys: Final = LoggedLiteLLMParams.__annotations__.keys() - for k, v in request_data.items(): + for k, v in _proxy_only_error_log_fields(request_fields, logged_metadata_keys).items(): if k in litellm_param_keys: _litellm_params[k] = v elif k not in ("model", "user", "litellm_logging_obj"): _optional_params[k] = v + existing: Final = _as_str_mapping( + litellm_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + litellm_params: Final = { + **_litellm_params, + "metadata": { + **_as_str_mapping(existing.get("metadata")), + **_as_str_mapping(_litellm_params.get("metadata")), + **user_api_key_logged_metadata, + }, + **{ + key: { + **_as_str_mapping(existing.get(key)), + **_as_str_mapping(value), + **user_api_key_logged_metadata, + } + for key, value in _litellm_params.items() + if key == "litellm_metadata" + }, + } + attribution: Final = _stamp_deployment_attribution( - _litellm_params, + litellm_params, request_data.get("model"), user_api_key_dict.team_id, dispatched=_reached_deployment(litellm_logging_obj), @@ -3923,7 +4011,7 @@ class ProxyLogging: model=request_data.get("model", ""), user=request_data.get("user", ""), optional_params=_optional_params, - litellm_params=_litellm_params, + litellm_params=litellm_params, **( {"custom_llm_provider": attribution["custom_llm_provider"]} if "custom_llm_provider" in attribution diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py index 07874bdbeb5..7ab683e9618 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -120,6 +120,24 @@ class TestGCSBucketBase: custom_llm_provider="vertex_ai", ) + @pytest.mark.asyncio + async def test_gcs_logging_config_ignores_request_time_bucket_and_credentials(self, monkeypatch): + monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/admin/sa.json") + handler = GCSBucketBase(bucket_name="admin-bucket") + + with patch.object(handler, "get_or_create_vertex_instance", new_callable=AsyncMock) as mock_vertex: + config = await handler.get_gcs_logging_config( + { + "standard_callback_dynamic_params": { + "gcs_bucket_name": "caller-bucket", + "gcs_path_service_account": "/caller/sa.json", + } + } + ) + + assert (config["bucket_name"], config["path_service_account"]) == ("admin-bucket", "/admin/sa.json") + mock_vertex.assert_awaited_once_with(credentials="/admin/sa.json") + @pytest.mark.asyncio async def test_log_json_data_on_gcs_url_encodes_object_name(self): handler = GCSBucketBase(bucket_name="test-bucket") diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index a1f674b467f..e786753e1a8 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -3520,6 +3520,21 @@ def test_get_request_tags_from_metadata_and_litellm_metadata(): assert "User-Agent: litellm/1.0.0" in tags +@pytest.mark.parametrize("unusable_metadata_tags", [123, "tag", {"tag": True}, []]) +def test_get_request_tags_falls_back_to_litellm_metadata_when_metadata_tags_are_unusable(unusable_metadata_tags): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + tags = StandardLoggingPayloadSetup._get_request_tags( + litellm_params={ + "metadata": {"tags": unusable_metadata_tags}, + "litellm_metadata": {"tags": ["litellm-metadata-tag"]}, + }, + proxy_server_request={}, + ) + + assert tags == ["litellm-metadata-tag"] + + def test_get_request_tags_does_not_mutate_original_tags(): """ Test that _get_request_tags does not mutate the original tags list in metadata. diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 9b7cd127b83..39310836bf2 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -260,87 +260,3 @@ class TestVertexAIFilesHandler: ) assert service_account == '{"type": "service_account", "project_id": "p"}' - - @pytest.mark.asyncio - async def test_afile_content_honors_per_model_bucket_over_env(self, monkeypatch): - """ - Regression for #32640: a batch output written to a per-model gcs_bucket_name must be - readable even when the global GCS_BUCKET_NAME points at a different bucket. Before the - fix the read path resolved the bucket from env only and raised - "file_id bucket does not match the configured storage bucket". - """ - monkeypatch.setenv("GCS_BUCKET_NAME", "env-default-bucket") - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - - file_id = "gs%3A%2F%2Fmy-model-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-batch-output.jsonl" - file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None) - - with ( - patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, - patch.object( - self.handler, - "get_or_create_vertex_instance", - new_callable=AsyncMock, - return_value=object(), - ), - ): - mock_download.return_value = b"batch output" - - result = await self.handler.afile_content( - file_content_request=file_content_request, - vertex_credentials="/model/sa.json", - vertex_project="test-project", - vertex_location="us-central1", - timeout=60.0, - max_retries=0, - litellm_params={ - "gcs_bucket_name": "my-model-bucket", - "vertex_credentials": "/model/sa.json", - }, - ) - - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == b"batch output" - - dynamic_params = mock_download.call_args.kwargs["standard_callback_dynamic_params"] - assert dynamic_params["gcs_bucket_name"] == "my-model-bucket" - assert dynamic_params["gcs_path_service_account"] == "/model/sa.json" - assert mock_download.call_args.kwargs["object_name"] == "litellm-vertex-files/uploads/abc-batch-output.jsonl" - - @pytest.mark.asyncio - async def test_afile_content_reads_without_global_env_bucket(self, monkeypatch): - """ - Regression for #32640: with no global GCS_BUCKET_NAME set, a model-group-level - deployment (per-model gcs_bucket_name) must still be readable. Before the fix the read - path raised "GCS_BUCKET_NAME is not set in the environment". - """ - monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - - file_id = "gs%3A%2F%2Fmy-model-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-batch-output.jsonl" - file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None) - - with ( - patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, - patch.object( - self.handler, - "get_or_create_vertex_instance", - new_callable=AsyncMock, - return_value=object(), - ), - ): - mock_download.return_value = b"batch output" - - result = await self.handler.afile_content( - file_content_request=file_content_request, - vertex_credentials="/model/sa.json", - vertex_project="test-project", - vertex_location="us-central1", - timeout=60.0, - max_retries=0, - litellm_params={"gcs_bucket_name": "my-model-bucket"}, - ) - - assert isinstance(result, HttpxBinaryResponseContent) - dynamic_params = mock_download.call_args.kwargs["standard_callback_dynamic_params"] - assert dynamic_params["gcs_bucket_name"] == "my-model-bucket" diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index db0a74d360a..112e4f6cb6c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1922,6 +1922,40 @@ async def test_pass_through_request_uses_resolved_timeout(): assert mock_get_client.call_args[1]["params"]["timeout"] == 1500 +@pytest.mark.asyncio +@pytest.mark.parametrize("allow_client_pricing_override", [True, False]) +async def test_pass_through_request_strips_client_guardrail_information(allow_client_pricing_override: bool): + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=RuntimeError("stop before upstream")) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + with pytest.raises(ProxyException): + await pass_through_request( + request=mock_request, + target="http://test.com", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth( + metadata={"allow_client_pricing_override": allow_client_pricing_override} + ), + custom_body={ + "model": "gpt-4o", + "input_cost_per_token": 0.5, + "metadata": { + "trace_id": "keep-me", + "standard_logging_guardrail_information": [{"guardrail_name": "forged", "guardrail_cost": 5.0}], + }, + }, + ) + + hooked_body: Final = mock_proxy_logging.pre_call_hook.call_args.kwargs["data"] + assert hooked_body["metadata"] == {"trace_id": "keep-me"} + assert ("input_cost_per_token" in hooked_body) is allow_client_pricing_override + + @pytest.mark.asyncio async def test_create_pass_through_route_forwards_timeout(): unique_path = "/test/path/unique/timeout" diff --git a/tests/unit/proxy/test_pricing_field_strip.py b/tests/unit/proxy/test_pricing_field_strip.py index 3e9bee9217d..c4a14b39f54 100644 --- a/tests/unit/proxy/test_pricing_field_strip.py +++ b/tests/unit/proxy/test_pricing_field_strip.py @@ -22,6 +22,7 @@ from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, _strip_client_pricing_overrides, add_litellm_data_to_request, + strip_client_guardrail_information, ) from litellm.types.utils import CustomPricingLiteLLMParams @@ -133,7 +134,7 @@ class TestStripClientPricingOverrides: "standard_logging_guardrail_information": [{"guardrail_cost": 5.0}], }, } - _strip_client_pricing_overrides(data) + strip_client_guardrail_information(data) assert data["metadata"] == {"user_session": "keep-me"} assert data["litellm_metadata"] == {} @@ -164,7 +165,6 @@ class TestStripClientPricingOverrides: def test_metadata_field_set_contains_model_info(self): assert "model_info" in _CLIENT_PRICING_METADATA_FIELDS - assert "standard_logging_guardrail_information" in _CLIENT_PRICING_METADATA_FIELDS def test_strip_emits_debug_log_listing_dropped_fields(self, caplog): # Operators need a paper trail so they can diagnose why a previously @@ -315,6 +315,33 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in(): assert updated["base_model"] == "gpt-4o-mini" +@pytest.mark.asyncio +@pytest.mark.parametrize("opt_in_field", ["metadata", "team_metadata"]) +async def test_add_litellm_data_to_request_strips_guardrail_information_despite_pricing_opt_in(opt_in_field): + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "input_cost_per_token": 0.0001, + "metadata": { + "user_session": "keep-me", + "standard_logging_guardrail_information": [{"guardrail_name": "forged", "guardrail_cost": 5.0}], + }, + } + + updated = await add_litellm_data_to_request( + data=data, + request=_make_request_mock(), + user_api_key_dict=_user_api_key_auth(**{opt_in_field: {"allow_client_pricing_override": True}}), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["input_cost_per_token"] == 0.0001 + assert updated["metadata"]["user_session"] == "keep-me" + assert "standard_logging_guardrail_information" not in updated["metadata"] + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_strips_json_string_litellm_metadata(): """``litellm_metadata`` may arrive as a JSON-encoded string (multipart/ diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index b5940e4c212..bc13595df6c 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -127,51 +127,132 @@ async def test_proxy_only_error_log_marks_no_upstream_llm_call(): @pytest.mark.asyncio -async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): - """Responses API requests carry guardrail info under ``litellm_metadata`` - (not ``metadata``). It must land in litellm_params so - ``merge_litellm_metadata`` can surface ``guardrail_information`` in the - spend-log failure row, matching the chat completions path.""" - from litellm.proxy._types import UserAPIKeyAuth +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_proxy_only_error_log_drops_caller_guardrail_information_without_logging_object( + metadata_field: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + guardrail_info: Final = [ + { + "guardrail_name": "caller-guardrail", + "guardrail_status": "guardrail_intervened", + "guardrail_cost": 123.45, + } + ] + request_data: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + metadata_field: {"tags": ["t"], "standard_logging_guardrail_information": guardrail_info}, + } - proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) - captured = {} - guardrail_info = [{"guardrail_name": "test-guard", "guardrail_status": "blocked"}] + await ProxyLogging(user_api_key_cache=DualCache())._handle_logging_proxy_only_error( + request_data=request_data, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-rejected", team_id="team-a", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=HTTPException(status_code=400, detail="blocked"), + ) - def fake_update_environment_variables(self, *args, **kwargs): - captured["litellm_params"] = kwargs.get("litellm_params") - captured["optional_params"] = kwargs.get("optional_params") + details: Final = request_data["litellm_logging_obj"].model_call_details + assert details["standard_logging_object"] is not None + assert details["litellm_params"][metadata_field]["tags"] == ["t"] + assert "standard_logging_guardrail_information" not in details["litellm_params"][metadata_field] + assert "guardrail_cost" not in str(details["standard_logging_object"]) + assert "litellm_metadata" not in details["optional_params"] - from litellm.litellm_core_utils.litellm_logging import Logging - orig_update_env = Logging.update_environment_variables - orig_pre_call = Logging.pre_call - orig_async_failure = Logging.async_failure_handler +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_proxy_only_error_log_keeps_tags_and_drops_caller_logging_fields( + metadata_field: str, +) -> None: + request_data: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + metadata_field: { + "langsmith_base_url": "https://caller.invalid", + "langsmith_api_key": "caller-key", + "tags": ["t"], + "user_api_key_team_id": "caller-team", + "user_api_key_hash": "caller-hash", + "user_api_key_auth_metadata": {"caller": "metadata"}, + }, + "langsmith_base_url": "https://caller.invalid", + "standard_callback_dynamic_params": {"gcs_path_service_account": "/caller/sa.json"}, + } - async def _noop_async_failure(self, *args, **kwargs): - return None + await ProxyLogging(user_api_key_cache=DualCache())._handle_logging_proxy_only_error( + request_data=request_data, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-rejected", team_id="team-a", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=Exception("bad key"), + ) - Logging.update_environment_variables = fake_update_environment_variables - Logging.pre_call = lambda self, *args, **kwargs: None - Logging.async_failure_handler = _noop_async_failure - try: - await proxy_logging_obj._handle_logging_proxy_only_error( - request_data={ - "model": "gpt-4o", - "input": "blocked prompt", - "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, - }, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/responses"), - route="/v1/responses", - original_exception=HTTPException(status_code=400, detail="blocked"), - ) - finally: - Logging.update_environment_variables = orig_update_env - Logging.pre_call = orig_pre_call - Logging.async_failure_handler = orig_async_failure + details: Final = request_data["litellm_logging_obj"].model_call_details + assert "caller" not in str(details["standard_callback_dynamic_params"]) + assert "langsmith_base_url" not in details + assert "caller" not in str(details["litellm_params"]) + assert details["litellm_params"][metadata_field]["tags"] == ["t"] + assert details["litellm_params"]["metadata"]["user_api_key_team_id"] == "team-a" + assert details["litellm_params"][metadata_field]["user_api_key_team_id"] == "team-a" - assert captured["litellm_params"]["litellm_metadata"]["standard_logging_guardrail_information"] == guardrail_info - assert "litellm_metadata" not in captured["optional_params"] + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_proxy_only_error_log_drops_malformed_metadata_tags(metadata_field: str) -> None: + guardrail_info: Final = [{"guardrail_name": "g", "guardrail_status": "guardrail_intervened"}] + request_data: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + metadata_field: {"tags": 123, "standard_logging_guardrail_information": guardrail_info}, + } + + await ProxyLogging(user_api_key_cache=DualCache())._handle_logging_proxy_only_error( + request_data=request_data, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-rejected", team_id="team-a", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=Exception("bad key"), + ) + + details: Final = request_data["litellm_logging_obj"].model_call_details + assert details["standard_logging_object"] is not None + assert "tags" not in details["litellm_params"][metadata_field] + assert "standard_logging_guardrail_information" not in details["litellm_params"][metadata_field] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_proxy_only_error_log_builds_payload_when_logging_object_has_malformed_tags( + metadata_field: str, +) -> None: + request_data: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "litellm_call_id": "malformed-tags-call", + metadata_field: {"tags": 123}, + } + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=real_datetime.datetime.now(), + **request_data, + ) + + await ProxyLogging(user_api_key_cache=DualCache())._handle_logging_proxy_only_error( + request_data={**request_data, "litellm_logging_obj": logging_obj}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-blocked", team_id="team-a", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=HTTPException(status_code=400, detail="Content blocked"), + ) + + standard_logging_object: Final = logging_obj.model_call_details["standard_logging_object"] + assert standard_logging_object is not None + assert standard_logging_object["model"] == request_data["model"] + assert standard_logging_object["metadata"]["user_api_key_team_id"] == "team-a" + assert 123 not in standard_logging_object["request_tags"] @pytest.mark.asyncio diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 6557f0d1446..f880d3548f7 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -402,6 +402,103 @@ async def test_post_call_failure_hook_flags_pre_routing_reject_despite_caller_mo assert kwargs["litellm_params"]["custom_llm_provider"] == "openai" +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_pre_routing_reject_logs_the_configured_deployment_over_caller_metadata( + proxy_logging, make_user_api_key_auth, monkeypatch, metadata_field: str +): + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + recorded: list[dict] = [] + + class _RecordingLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + recorded.append(kwargs) + + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "internal-model", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"}, + } + ] + ), + ) + monkeypatch.setattr(litellm, "callbacks", [_RecordingLogger()]) + proxy_logging.alert_types = [] + request_data: Final[dict[str, object]] = { + "model": "internal-model", + "messages": [{"role": "user", "content": "hi"}], + "litellm_call_id": "caller-routing-metadata", + metadata_field: {"deployment": "openai/caller-chosen", "model_group": "caller-group"}, + } + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", rules_obj=litellm.utils.Rules(), start_time=datetime.now(), **request_data + ) + + await proxy_logging.post_call_failure_hook( + request_data={**request_data, "litellm_logging_obj": logging_obj}, + original_exception=HTTPException(status_code=400, detail="Content blocked"), + user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"), + route="/chat/completions", + ) + + assert len(recorded) == 1 + now: Final = datetime.now() + payload: Final = get_logging_payload( + kwargs={**recorded[0], "completion_start_time": now}, response_obj=None, start_time=now, end_time=now + ) + standard_logging_object: Final = recorded[0]["standard_logging_object"] + assert { + "logged_model": standard_logging_object["model"], + "logged_model_group": standard_logging_object["model_group"], + "spend_model": payload["model"], + "spend_model_group": payload["model_group"], + } == { + "logged_model": "openai/gpt-4.1", + "logged_model_group": "internal-model", + "spend_model": "openai/gpt-4.1", + "spend_model_group": "internal-model", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("error_type", "logged_tags"), [(ProxyErrorTypes.auth_error, []), (None, ["checkout"])]) +async def test_proxy_only_failure_log_keeps_caller_tags_only_after_auth( + proxy_logging, monkeypatch, error_type: ProxyErrorTypes | None, logged_tags: list[str] +): + from litellm.proxy import proxy_server + + recorded: list[dict] = [] + + class _RecordingLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + recorded.append(kwargs) + + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [_RecordingLogger()]) + proxy_logging.alert_types = [] + + await proxy_logging.post_call_failure_hook( + request_data={ + "model": "internal-model", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["checkout"]}, + }, + original_exception=HTTPException(status_code=401, detail="Invalid proxy server token passed"), + user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"), + error_type=error_type, + route="/chat/completions", + ) + + assert len(recorded) == 1 + assert recorded[0]["standard_logging_object"]["request_tags"] == logged_tags + + @pytest.mark.asyncio async def test_post_call_failure_hook_attribution_does_not_count_against_the_deployment( proxy_logging, make_user_api_key_auth, monkeypatch @@ -574,9 +671,8 @@ async def test_post_call_failure_hook_omits_provider_when_a_deployment_does_not_ async def test_handle_logging_proxy_only_path_attributes_with_read_only_metadata( proxy_logging, make_user_api_key_auth, monkeypatch ): - """With a logging object already on the request, its metadata is taken as given; - a read-only mapping there must not crash the stamp, and the failure handler - still receives the provider attribution.""" + """With a logging object already on the request, allowlisted metadata is kept, + a read-only mapping must not crash the stamp, and the failure handler receives provider attribution.""" from litellm.proxy import proxy_server monkeypatch.setattr( @@ -595,16 +691,26 @@ async def test_handle_logging_proxy_only_path_attributes_with_read_only_metadata logging_obj = MagicMock() logging_obj.call_type = "acompletion" logging_obj.model_call_details = {} + logging_obj.litellm_params = {"metadata": {"preexisting": "logged"}} logging_obj.async_failure_handler = AsyncMock() + guardrail_info: Final = [ + { + "guardrail_name": "local-block", + "guardrail_status": "guardrail_intervened", + "guardrail_cost": 0.25, + } + ] await proxy_logging._handle_logging_proxy_only_error( request_data={ "litellm_logging_obj": logging_obj, "model": "internal-model", "messages": [{"role": "user", "content": "hi"}], - "metadata": MappingProxyType({"user_api_key_alias": "frozen"}), + "metadata": MappingProxyType( + {"tags": ["checkout"], "standard_logging_guardrail_information": guardrail_info} + ), }, - user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"), + user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions", key_alias="frozen"), route="/chat/completions", original_exception=HTTPException(status_code=403, detail="blocked"), ) @@ -613,7 +719,11 @@ async def test_handle_logging_proxy_only_path_attributes_with_read_only_metadata update_kwargs = logging_obj.update_environment_variables.call_args.kwargs assert update_kwargs["custom_llm_provider"] == "openai" assert update_kwargs["litellm_params"]["custom_llm_provider"] == "openai" - assert update_kwargs["litellm_params"]["metadata"] == {"user_api_key_alias": "frozen"} + assert { + key: update_kwargs["litellm_params"]["metadata"][key] for key in ("preexisting", "user_api_key_alias") + } == {"preexisting": "logged", "user_api_key_alias": "frozen"} + assert update_kwargs["litellm_params"]["metadata"]["tags"] == ["checkout"] + assert update_kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] == guardrail_info @pytest.mark.asyncio From b766aaf010442204c3770b72bccdb87b1ea97479 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 12:43:04 -0700 Subject: [PATCH 009/147] fix(ui): expose the authorization MCP auth type in create and edit forms (#48629) * test(ui): cover the authorization MCP auth type in create and edit forms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): expose the authorization MCP auth type in create and edit forms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../CreateMCPServer.integration.test.tsx | 40 +++++++++++++++++++ .../_components/createServerPayload.ts | 1 + .../editServerPayload.differential.cases.ts | 1 + .../mcp_server_edit.integration.test.tsx | 5 +++ .../_components/mcp_server_edit.tsx | 2 +- .../src/components/mcp_tools/types.tsx | 2 + 6 files changed, 50 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx index 65656a053ce..3b3269b87cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx @@ -396,6 +396,46 @@ describe("CreateMCPServer", () => { expect(payload.credentials).toEqual({ auth_value: "my-secret-key" }); }); + it("creates an Authorization (raw header) server with the header value sent verbatim", async () => { + await selectHttpTransport(); + + fireEvent.change(getServerNameInput(), { target: { value: "Raw_Auth_Server" } }); + fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), { + target: { value: "https://example.com/mcp" }, + }); + + await selectOption("Authentication", "Authorization (raw header)"); + + fireEvent.change(await screen.findByPlaceholderText("Enter token or secret"), { + target: { value: "ApiKey raw-key-123" }, + }); + + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-2", + server_name: "Raw_Auth_Server", + alias: "Raw_Auth_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "authorization", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); + + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: "Add MCP Server" })); + }); + + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.auth_type).toBe("authorization"); + expect(payload.credentials).toEqual({ auth_value: "ApiKey raw-key-123" }); + }); + it("does not write the browser-authorized token into form.credentials for true_passthrough", async () => { await selectHttpTransport(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts index 30c4c502c42..5536b4b1bc4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts @@ -15,6 +15,7 @@ export const AUTH_TYPES_REQUIRING_AUTH_VALUE = [ AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC, + AUTH_TYPE.AUTHORIZATION, ]; export const AUTH_TYPES_REQUIRING_CREDENTIALS = [ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.cases.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.cases.ts index 4219113918b..e9f32b5c4db 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.cases.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.cases.ts @@ -82,6 +82,7 @@ const AUTH_TYPES = [ "bearer_token", "token", "basic", + "authorization", "oauth2", "oauth2_token_exchange", "oauth2_id_jag", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx index 9f91c06e9e5..15b09961a84 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx @@ -135,6 +135,11 @@ const CASES: readonly Case[] = [ server: { ...BASE, auth_type: "basic", credentials: CREDS }, expected: { ...EXPECTED_BASE, auth_type: "basic", credentials: { auth_value: "secret-value" } }, }, + { + label: "authorization sends only auth_value under credentials", + server: { ...BASE, auth_type: "authorization", credentials: CREDS }, + expected: { ...EXPECTED_BASE, auth_type: "authorization", credentials: { auth_value: "secret-value" } }, + }, { label: "oauth2 m2m mounts token_url and omits the interactive endpoints entirely", server: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index f35e3ab3d3a..57985fd8cbb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -51,6 +51,7 @@ import OAuthFormFields from "./OAuthFormFields"; import MCPLogoSelector from "./MCPLogoSelector"; import EnvVarsSection from "./EnvVarsSection"; import { validateMCPServerUrl, validateMCPServerName, normalizeToolOverrideMap } from "./utils"; +import { AUTH_TYPES_REQUIRING_AUTH_VALUE } from "./createServerPayload"; import { EditServerFormValues, buildEditServerPayload, editPayloadErrorMessage } from "./editServerPayload"; import { DUPLICATE_IDENTIFIER_MESSAGE, findDuplicateMcpServer, mcpSubmitErrorReason } from "./duplicateServerCheck"; import { toast } from "@/lib/toast"; @@ -94,7 +95,6 @@ interface MCPServerEditProps { stdioEnabled?: boolean; } -const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC]; export const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; const MCPServerEdit: React.FC = ({ diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 22fd64f10d8..61c81c030c5 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -38,6 +38,7 @@ export const AUTH_TYPE = { BEARER_TOKEN: "bearer_token", TOKEN: "token", BASIC: "basic", + AUTHORIZATION: "authorization", OAUTH2: "oauth2", OAUTH2_TOKEN_EXCHANGE: "oauth2_token_exchange", OAUTH2_ID_JAG: "oauth2_id_jag", @@ -52,6 +53,7 @@ export const AUTH_TYPE_ITEMS = [ { value: AUTH_TYPE.BEARER_TOKEN, label: "Bearer Token" }, { value: AUTH_TYPE.TOKEN, label: "Token" }, { value: AUTH_TYPE.BASIC, label: "Basic Auth" }, + { value: AUTH_TYPE.AUTHORIZATION, label: "Authorization (raw header)" }, { value: AUTH_TYPE.OAUTH2, label: "OAuth" }, { value: AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE, label: "OAuth Token Exchange (OBO)" }, { value: AUTH_TYPE.OAUTH2_ID_JAG, label: "ID-JAG (Okta Cross App Access)" }, From 29946e4b75421d5a91e4e3e7f8df913a4f956032 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 12:43:36 -0700 Subject: [PATCH 010/147] feat(hosted_vllm): accept images on decision deployments (#48621) * test(decisions): cover hosted_vllm decision images Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(hosted_vllm): accept images on decision deployments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../hosted_vllm/decisions/transformation.py | 2 + tests/unit/decisions/test_main.py | 124 ++++++++++++++++++ 2 files changed, 126 insertions(+) diff --git a/litellm/llms/hosted_vllm/decisions/transformation.py b/litellm/llms/hosted_vllm/decisions/transformation.py index ef99f2de173..fab07d2e7db 100644 --- a/litellm/llms/hosted_vllm/decisions/transformation.py +++ b/litellm/llms/hosted_vllm/decisions/transformation.py @@ -9,6 +9,8 @@ class HostedVLLMDecisionsConfig(BaseDecisionsConfig): api_key_env = ("HOSTED_VLLM_API_KEY",) api_base_env = ("HOSTED_VLLM_API_BASE",) api_key_required = False + supports_images = True + inlines_remote_images = True health_check_questions = MappingProxyType( { "reachable": MappingProxyType( diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 0bc18080ef4..b9813289749 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -4026,3 +4026,127 @@ async def test_lm_studio_decisions_sends_inline_images_to_systemone( "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, } assert response.answers[0].probability == 0.9 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("remote_image", (False, True), ids=("inline", "remote")) +async def test_hosted_vllm_decisions_sends_input_images_to_systemone( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + remote_image: bool, + async_only_image_fetch: AsyncOnlyImageFetch, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + image: Final[str] = "https://images.example/input.png" if remote_image else "data:image/png;base64,AA==" + expected_image: Final[str] = async_only_image_fetch.data_url if remote_image else image + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond(json=_RESPONSE) + + response: Final = await litellm.adecisions( + model="hosted_vllm/Qwen/Qwen3-0.6B", + input=[ + { + "role": "user", + "content": [{"type": "input_text", "text": "review"}, {"type": "input_image", "image_url": image}], + } + ], + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + api_base="http://vllm.local:8000/v1", + ) + + assert json.loads(route.calls.last.request.content) == { + "model": "Qwen/Qwen3-0.6B", + "state": "review", + "images": [expected_image], + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert response.answers[0].probability == 0.9 + assert async_only_image_fetch.fetched == (["https://images.example/input.png"] if remote_image else []) + + +def test_hosted_vllm_systemone_forwards_images( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + image: Final[str] = "data:image/png;base64,AA==" + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond(json=_RESPONSE) + + litellm.systemone( + model="hosted_vllm/Qwen/Qwen3-0.6B", + state="review", + images=[image], + questions={ + "is_defect": { + "type": "choice", + "instructions": "Is this a defect?", + "criteria": {"yes": "defective", "no": "intact"}, + } + }, + api_base="http://vllm.local:8000", + ) + + assert json.loads(route.calls.last.request.content) == { + "model": "Qwen/Qwen3-0.6B", + "state": "review", + "images": [image], + "questions": { + "is_defect": { + "type": "choice", + "instructions": "Is this a defect?", + "criteria": {"yes": "defective", "no": "intact"}, + } + }, + } + + +@pytest.mark.asyncio +async def test_hosted_vllm_decisions_sends_image_only_input_to_systemone( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + image: Final[str] = "data:image/png;base64,AA==" + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond(json=_RESPONSE) + + await litellm.adecisions( + model="hosted_vllm/Qwen/Qwen3-0.6B", + input=[{"role": "user", "content": [{"type": "input_image", "image_url": image}]}], + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + api_base="http://vllm.local:8000", + ) + + assert json.loads(route.calls.last.request.content) == { + "model": "Qwen/Qwen3-0.6B", + "state": "", + "images": [image], + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + + +def test_hosted_vllm_systemone_maps_image_rejection_to_bad_request( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + message: Final[str] = "this model does not take image inputs" + route: Final = respx_mock.post("http://vllm.local:8000/v1/systemone").respond( + status_code=400, + json={"error": {"message": message, "type": "BadRequestError", "code": 400}}, + ) + + with pytest.raises(litellm.BadRequestError, match=re.escape(message)) as caught: + litellm.systemone( + model="hosted_vllm/Qwen/Qwen3-0.6B", + state="review", + images=["data:image/png;base64,AA=="], + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + api_base="http://vllm.local:8000", + num_retries=0, + ) + + assert caught.value.status_code == 400 + assert route.call_count == 1 From 618d5aec71880362a78984a839e080e403ec99d3 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:46:19 -0700 Subject: [PATCH 011/147] fix(mcp): load deprecated package aliases on demand (#48615) * fix(mcp): load deprecated package aliases on demand * fix(mcp): preserve lazy child access and module reloads --- .../litellm_core_utils/deprecated_package.py | 84 ++++++++++++------- .../unit/experimental_mcp_client/test_init.py | 2 + .../_experimental/mcp_server/test_init.py | 49 +++++++++++ 3 files changed, 103 insertions(+), 32 deletions(-) diff --git a/litellm/litellm_core_utils/deprecated_package.py b/litellm/litellm_core_utils/deprecated_package.py index 542686ecf19..f9915af4880 100644 --- a/litellm/litellm_core_utils/deprecated_package.py +++ b/litellm/litellm_core_utils/deprecated_package.py @@ -1,53 +1,73 @@ import importlib +import importlib.abc +import importlib.machinery +import importlib.util import sys import warnings -from itertools import chain -from pathlib import Path +from collections.abc import Sequence +from dataclasses import dataclass +from functools import partial from types import ModuleType from typing import Final -def _dotted_name(root: Path, path: Path, package_name: str) -> str: - parts: Final = path.relative_to(root).with_suffix("").parts - return ".".join((package_name, *parts)).removesuffix(".__init__") +@dataclass(frozen=True, slots=True) +class _AliasLoader(importlib.abc.Loader): + new_dotted: str + canonical_spec: importlib.machinery.ModuleSpec + old_dotted: str + + def create_module(self, spec: importlib.machinery.ModuleSpec) -> ModuleType: + return importlib.import_module(self.new_dotted) + + def exec_module(self, module: ModuleType) -> None: + module.__spec__ = self.canonical_spec # rebind-ok: [LIT011] restore import metadata overwritten by importlib + if self.canonical_spec.submodule_search_locations is not None: + _alias_attributes(self.old_dotted, self.new_dotted) -def _parents_and_self(dotted: str, package_name: str) -> tuple[str, ...]: - parts: Final = dotted.removeprefix(f"{package_name}.").split(".") - return tuple(f"{package_name}.{'.'.join(parts[:depth])}" for depth in range(1, len(parts) + 1)) +@dataclass(frozen=True, slots=True) +class _AliasFinder(importlib.abc.MetaPathFinder): + old_name: str + new_name: str - -def _module_names(package: ModuleType) -> tuple[str, ...]: - roots: Final = tuple(Path(root) for root in package.__path__) - files: Final = tuple( - (root, path) for root in roots for path in root.rglob("*.py") - ) # comprehension-ok: nested walk over each package root - leaves: Final = tuple( - name - for name in (_dotted_name(root, path, package.__name__) for root, path in files) - if name != package.__name__ - ) - return tuple( - sorted( - set(chain.from_iterable(_parents_and_self(leaf, package.__name__) for leaf in leaves)) - {package.__name__} + def find_spec( + self, + fullname: str, + path: Sequence[str] | None = None, + target: ModuleType | None = None, + ) -> importlib.machinery.ModuleSpec | None: + if not fullname.startswith(f"{self.old_name}."): + return None + new_dotted: Final = f"{self.new_name}{fullname.removeprefix(self.old_name)}" + new_spec: Final = importlib.util.find_spec(new_dotted) + if new_spec is None: + return None + return importlib.util.spec_from_loader( + fullname, + _AliasLoader(new_dotted, new_spec, fullname), + is_package=new_spec.submodule_search_locations is not None, ) - ) -def _alias(old_name: str, new_package_name: str, new_dotted: str) -> None: - module: Final = importlib.import_module(new_dotted) - old_dotted: Final = f"{old_name}{new_dotted.removeprefix(new_package_name)}" - old_parent, _, child = old_dotted.rpartition(".") - sys.modules[old_dotted] = module - setattr(sys.modules[old_parent], child, module) +def _deprecated_attribute(old_name: str, new_name: str, name: str) -> ModuleType: + if importlib.util.find_spec(f"{new_name}.{name}") is None: + raise AttributeError(f"module {old_name!r} has no attribute {name!r}") + return importlib.import_module(f"{old_name}.{name}") + + +def _alias_attributes(old_name: str, new_name: str) -> None: + if "__getattr__" not in vars(sys.modules[old_name]): + sys.modules[old_name].__getattr__ = partial(_deprecated_attribute, old_name, new_name) def alias_deprecated_package(old_name: str, new_package: ModuleType) -> None: - """Make every module under ``new_package`` importable and attribute-reachable under ``old_name`` and warn once.""" warnings.warn( f"{old_name} is deprecated, import {new_package.__name__} instead", DeprecationWarning, stacklevel=3, ) - for new_dotted in _module_names(new_package): - _alias(old_name, new_package.__name__, new_dotted) + finder: Final = _AliasFinder(old_name, new_package.__name__) + if finder not in sys.meta_path: + sys.meta_path.insert(0, finder) + _alias_attributes(old_name, new_package.__name__) diff --git a/tests/unit/experimental_mcp_client/test_init.py b/tests/unit/experimental_mcp_client/test_init.py index 9d2f1137470..1ce2a236d9d 100644 --- a/tests/unit/experimental_mcp_client/test_init.py +++ b/tests/unit/experimental_mcp_client/test_init.py @@ -19,6 +19,8 @@ def test_old_package_warns_and_deep_imports_resolve_to_mcp_client(monkeypatch: p with pytest.warns(DeprecationWarning, match="litellm.mcp_client"): import litellm.experimental_mcp_client as old + assert old.client is litellm.mcp_client.client + assert old.tools is litellm.mcp_client.tools assert old.__all__ == litellm.mcp_client.__all__ assert old.load_mcp_tools is litellm.mcp_client.load_mcp_tools assert old.call_openai_tool is litellm.mcp_client.call_openai_tool diff --git a/tests/unit/proxy/_experimental/mcp_server/test_init.py b/tests/unit/proxy/_experimental/mcp_server/test_init.py index 2f933e11351..5ecb79b9813 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_init.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_init.py @@ -1,6 +1,7 @@ import importlib import sys from collections.abc import Mapping +from typing import Final import pytest @@ -16,9 +17,17 @@ def test_old_package_warns_and_deep_imports_resolve_to_proxy_mcp_server(monkeypa if name.startswith(OLD_PACKAGE): monkeypatch.delitem(sys.modules, name) + monkeypatch.delattr(litellm.proxy.mcp_server.auth, "user_api_key_auth_mcp") + monkeypatch.delitem(sys.modules, "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp") + with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"): old = importlib.import_module(OLD_PACKAGE) + assert old.server is litellm.proxy.mcp_server.server + assert old.auth is litellm.proxy.mcp_server.auth + assert old.auth.user_api_key_auth_mcp is importlib.import_module( + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp" + ) assert importlib.import_module(f"{OLD_PACKAGE}.server") is litellm.proxy.mcp_server.server assert importlib.import_module(f"{OLD_PACKAGE}.utils") is litellm.proxy.mcp_server.utils assert ( @@ -28,3 +37,43 @@ def test_old_package_warns_and_deep_imports_resolve_to_proxy_mcp_server(monkeypa attrs: Mapping[str, object] = vars(old) assert attrs["server"] is litellm.proxy.mcp_server.server assert attrs["auth"] is litellm.proxy.mcp_server.auth + + +def test_old_stdio_helper_import_works_without_optional_mcp_dependency(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delattr(litellm.proxy, "mcp_server") + old_parent: Final = importlib.import_module("litellm.proxy._experimental") + monkeypatch.delattr(old_parent, "mcp_server", raising=False) + monkeypatch.delattr(litellm, "mcp_client", raising=False) + prefixes: Final = (OLD_PACKAGE, "litellm.proxy.mcp_server", "litellm.mcp_client", "mcp") + for name in tuple(sys.modules): + if any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes): + monkeypatch.delitem(sys.modules, name) + monkeypatch.setitem(sys.modules, "mcp", None) + + with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"): + old: Final = importlib.import_module(f"{OLD_PACKAGE}.stdio_gate") + new: Final = importlib.import_module("litellm.proxy.mcp_server.stdio_gate") + + monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true") + assert old.is_mcp_stdio_enabled() is True + monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "false") + assert old.is_mcp_stdio_enabled() is False + assert old is new + + +def test_legacy_child_import_keeps_canonical_module_reload_working(monkeypatch: pytest.MonkeyPatch) -> None: + for name in tuple(sys.modules): + if name.startswith(OLD_PACKAGE): + monkeypatch.delitem(sys.modules, name) + + helper: Final = importlib.import_module("litellm.proxy.mcp_server.stdio_gate") + with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"): + old: Final = importlib.import_module(f"{OLD_PACKAGE}.stdio_gate") + + monkeypatch.setattr(helper, "MCP_STDIO_ENABLED_ENV_VAR", "RELOAD_PROBE") + monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true") + monkeypatch.setenv("RELOAD_PROBE", "false") + importlib.reload(helper) + + assert old is helper + assert old.is_mcp_stdio_enabled() is True From 34134d7ae5b7909c6b75aedb0efc6eff9aab281f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 27 Oct 2026 12:46:22 -0700 Subject: [PATCH 012/147] fix(auth): log routine 401 rejections as one warning without a traceback (#48613) A missing or malformed Authorization header, a blocked key, a non-admin key on an admin-only or master-key-only route and a SCIM-deactivated user were raised as bare exceptions, so the auth failure handler could not see the 401 and logged each one at ERROR with a full traceback, which flooded the log under scanner traffic They now raise a typed 401 carrying the quiet-log marker that the non-sk bearer rejection already used, and the unknown sk- key rejection joins the same path, so each lands as one WARNING line on stdout with no traceback. The client response is unchanged --- litellm/proxy/auth/auth_exception_handler.py | 7 +-- litellm/proxy/auth/auth_utils.py | 20 +++++---- litellm/proxy/auth/route_checks.py | 14 ++++-- litellm/proxy/auth/user_api_key_auth.py | 45 ++++++++++++++++--- .../proxy/auth/test_auth_exception_handler.py | 29 ++++++++---- tests/unit/proxy/auth/test_route_checks.py | 2 +- .../unit/proxy/auth/test_user_api_key_auth.py | 4 +- 7 files changed, 88 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 2eb25cc6105..ed743764fce 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -24,6 +24,7 @@ from litellm.proxy.auth.auth_utils import ( # noqa: F401 # legacy module expor _get_request_ip_address, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export get_request_ip_address, is_invalid_virtual_key_error, + is_unauthorized_error, mark_invalid_virtual_key_error, normalize_request_route, ) @@ -200,11 +201,11 @@ class UserAPIKeyAuthExceptionHandler: user_agent: Final = _get_user_agent(request) # Log authentication failures before identity seeding and callbacks, so the log - # survives a raising callback pipeline. Classify and route malformed virtual-key - # rejections to WARNING on stdout (suppressible via LITELLM_LOG=ERROR). + # survives a raising callback pipeline. Route 401 rejections to WARNING on stdout + # (suppressible via LITELLM_LOG=ERROR); anything else is unexpected and logs at ERROR. log_extra: Final = {"requester_ip": requester_ip} is_invalid_virtual_key: Final = is_invalid_virtual_key_error(e) - is_quiet_log: Final = is_invalid_virtual_key and not litellm.log_client_error_tracebacks + is_quiet_log: Final = is_unauthorized_error(e) and not litellm.log_client_error_tracebacks logger: Final = verbose_proxy_stdout_logger if is_quiet_log else verbose_proxy_logger logger.log( logging.WARNING if is_quiet_log else logging.ERROR, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 30ede935cfa..aa6ddae438d 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -48,6 +48,16 @@ from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange from litellm.types.videos.main import DecodedVideoId +def is_unauthorized_error(exception: BaseException | None) -> bool: + """True when the proxy itself rejected the caller's credentials with a 401.""" + if not isinstance(exception, (HTTPException, ProxyException)): + return False + + code: Final[object] = getattr(exception, "code", None) + status_code: Final[object] = code if code is not None else getattr(exception, "status_code", None) + return str(status_code) == str(status.HTTP_401_UNAUTHORIZED) + + def is_invalid_virtual_key_error(exception: BaseException | None) -> bool: """True when an authentication error rejects a malformed virtual key. @@ -57,15 +67,7 @@ def is_invalid_virtual_key_error(exception: BaseException | None) -> bool: match would let a request body demote an authorization failure to the quiet log path. """ - if not isinstance(exception, (HTTPException, ProxyException)): - return False - - code: Final[object] = getattr(exception, "code", None) - status_code: Final[object] = code if code is not None else getattr(exception, "status_code", None) - if str(status_code) != str(status.HTTP_401_UNAUTHORIZED): - return False - - return getattr(exception, INVALID_VIRTUAL_KEY_ERROR_MARKER, False) is True + return is_unauthorized_error(exception) and getattr(exception, INVALID_VIRTUAL_KEY_ERROR_MARKER, False) is True def log_model_access_denial(exc: BaseException) -> None: diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index c402f111563..74e4c4c6d11 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -14,6 +14,8 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, + ProxyErrorTypes, + ProxyException, UserAPIKeyAuth, ) @@ -284,7 +286,7 @@ class RouteChecks: route (str): The route being accessed Raises: - Exception: With user role and masked user_id information + ProxyException: 401 with user role and masked user_id information """ user_role = "unknown" user_id = "unknown" @@ -293,8 +295,14 @@ class RouteChecks: user_id = user_obj.user_id or "unknown" masked_user_id: Final = RouteChecks._mask_user_id(user_id) - raise Exception( - f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" + raise ProxyException( + message=( + "Authentication Error, Only proxy admin can be used to generate, delete, update info for new " + f"keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" + ), + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_401_UNAUTHORIZED, ) @staticmethod diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b1e8bdbe787..cb490b55352 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2513,10 +2513,20 @@ async def user_api_key_auth_builder( parent_otel_span=parent_otel_span, ) elif api_key is None: # only require api key if master key is set - raise Exception("No api key passed in.") + raise ProxyException( + message="Authentication Error, No api key passed in.", + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) elif api_key == "": # missing 'Bearer ' prefix - raise Exception("Malformed API Key passed in. Ensure Key has `Bearer ` prefix.") + raise ProxyException( + message="Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix.", + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) if route == "/user/auth": if general_settings.get("allow_user_auth", False) is True: @@ -2752,7 +2762,12 @@ async def user_api_key_auth_builder( ## IF it's not a master key ## Route should not be in master_key_only_routes if route in LiteLLMRoutes.master_key_only_routes.value: - raise Exception(f"Tried to access route={route}, which is only for MASTER KEY") + raise ProxyException( + message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY", + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) ## Check DB @@ -2880,7 +2895,12 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ## base case ## key is disabled if valid_token.blocked is True: - raise Exception("Key is blocked. Update via `/key/unblock` if you're an admin.") + raise ProxyException( + message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) await enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, @@ -2923,8 +2943,14 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e and isinstance(user_obj.metadata, dict) and user_obj.metadata.get("scim_active") is False ): - raise Exception( - f"User={valid_token.user_id} has been deactivated via SCIM. Keys owned by this user cannot be used." + raise ProxyException( + message=( + f"Authentication Error, User={valid_token.user_id} has been deactivated via SCIM. " + "Keys owned by this user cannot be used." + ), + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_401_UNAUTHORIZED, ) # Check 2a. Check if model has zero cost - if so, skip all budget checks @@ -3241,7 +3267,12 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e if valid_token is None: # No token was found when looking up in the DB - raise Exception("Invalid proxy server token passed") + raise ProxyException( + message="Authentication Error, Invalid proxy server token passed", + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) if valid_token_dict is not None: virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj( user_obj=user_obj, diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index e21491fd996..92b4c93118f 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -668,7 +668,7 @@ async def test_expired_key_error_log_names_the_key_owner( ): verbose_proxy_logger.propagate = True try: - with caplog.at_level("ERROR", logger="LiteLLM Proxy"), pytest.raises(ProxyException): + with caplog.at_level("WARNING", logger="LiteLLM Proxy"), pytest.raises(ProxyException): await handler.handle_authentication_error( expired_key_error, MagicMock(), @@ -714,8 +714,8 @@ async def test_expired_key_error_log_names_the_key_owner( async def test_unknown_key_error_log_keeps_the_operator_wording_the_401_hides( looked_up, logged_fragment, absent_fragment, caplog ): - """A wrong key gets a generic 401, while the one auth failure log line keeps the - operator wording and the hash that was looked up, next to the requester IP.""" + """A wrong key gets a generic 401, while the one auth failure log line, a WARNING with no + traceback, keeps the operator wording and the hash that was looked up, next to the requester IP.""" from litellm.constants import INVALID_API_KEY_ERROR_MESSAGE from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError @@ -733,7 +733,7 @@ async def test_unknown_key_error_log_keeps_the_operator_wording_the_401_hides( ): verbose_proxy_logger.propagate = True try: - with caplog.at_level("ERROR", logger="LiteLLM Proxy"), pytest.raises(ProxyException) as exc_info: + with caplog.at_level("WARNING", logger="LiteLLM Proxy"), pytest.raises(ProxyException) as exc_info: await UserAPIKeyAuthExceptionHandler().handle_authentication_error( KeyNotFoundError(looked_up), MagicMock(), @@ -748,6 +748,7 @@ async def test_unknown_key_error_log_keeps_the_operator_wording_the_401_hides( assert exc_info.value.message == INVALID_API_KEY_ERROR_MESSAGE records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] assert len(records) == 1, [r.getMessage() for r in caplog.records] + assert records[0].levelname == "WARNING" and records[0].exc_info is None logged = records[0].getMessage() assert "Requester IP Address:" in logged and logged_fragment in logged, logged if absent_fragment is not None: @@ -1148,10 +1149,16 @@ def _marked_malformed_key_error() -> HTTPException: message="Authentication Error", type=ProxyErrorTypes.auth_error, param=None, code=401 ), False, - "ERROR", - id="expected_401_no_traceback", + "WARNING", + id="proxy_exception_401_warning_no_traceback", ), pytest.param(ValueError("unexpected internal error"), True, "ERROR", id="unexpected_error_keeps_traceback"), + pytest.param( + Exception("No api key passed in."), + True, + "ERROR", + id="bare_exception_coerced_to_401_still_keeps_traceback", + ), pytest.param( _marked_malformed_key_error(), False, @@ -1161,8 +1168,14 @@ def _marked_malformed_key_error() -> HTTPException: pytest.param( HTTPException(status_code=401, detail="LiteLLM Virtual Key expected. Received=test"), False, + "WARNING", + id="http_exception_401_warning_no_traceback", + ), + pytest.param( + ProxyException(message="Forbidden", type=ProxyErrorTypes.auth_error, param=None, code=403), + False, "ERROR", - id="phrase_without_marker_stays_loud", + id="403_stays_error_without_traceback", ), ], ) @@ -1189,7 +1202,7 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( try: try: raise auth_error - except (ProxyException, ValueError, HTTPException) as caught: + except Exception as caught: with caplog.at_level(expect_level, logger="LiteLLM Proxy"), pytest.raises((ProxyException, HTTPException)): await handler.handle_authentication_error( caught, diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 6de54aa4af6..b2c67b4e009 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -3742,7 +3742,7 @@ def test_team_spend_by_user_reachable_by_non_admin(user_role): return "allowed" assert outcome("/team/spend/by_user") == "allowed" - assert outcome("/team/spend/by_key").startswith("denied: Only proxy admin") + assert outcome("/team/spend/by_key").startswith("denied: Authentication Error, Only proxy admin") def test_user_daily_activity_aggregated_not_covered_by_prefix_match(): diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 60a8dd73a34..092805a9cac 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2031,9 +2031,9 @@ async def test_user_api_key_auth_websocket_rejection_adds_no_traceback_for_other assert rejection.value.status_code == 403 assert "expired key" in str(rejection.value.detail).lower() assert sent == [{"type": "websocket.close", "code": status.WS_1008_POLICY_VIOLATION, "reason": ""}] - proxy_records: Final = [record for record in caplog.records if record.name == verbose_proxy_logger.name] + proxy_records: Final = [record for record in caplog.records if record.name.startswith(verbose_proxy_logger.name)] assert [record.getMessage() for record in proxy_records if record.exc_info is not None] == [] - assert [record.levelno for record in proxy_records if record.levelno >= logging.WARNING] == [logging.ERROR] + assert [record.levelno for record in proxy_records if record.levelno >= logging.WARNING] == [logging.WARNING] @pytest.mark.asyncio From d65587593f04f6bb004f33e39c9f61d1fc542ed7 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 12:49:23 -0700 Subject: [PATCH 013/147] ci: spend-logs-guard requires kerry-berri approval on LiteLLM_SpendLogs schema and migration changes (#48518) * ci: spend-logs-guard requires kerry-berri approval on LiteLLM_SpendLogs schema and migration changes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(spend-logs-guard): flatten paginated reviews with jq, check out the base ref Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(spend-logs-guard): read the diff from the pulls/files API instead of fetching PR commits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Revert "ci(spend-logs-guard): read the diff from the pulls/files API instead of fetching PR commits" This reverts commit 36da4fe034426194431aa8af92f502b819dff981. --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/spend-logs-guard.yml | 49 ++++++++++ ci_cd/spend_logs_guard.py | 119 +++++++++++++++++++++++++ tests/unit/test_spend_logs_guard.py | 79 ++++++++++++++++ 3 files changed, 247 insertions(+) create mode 100644 .github/workflows/spend-logs-guard.yml create mode 100644 ci_cd/spend_logs_guard.py create mode 100644 tests/unit/test_spend_logs_guard.py diff --git a/.github/workflows/spend-logs-guard.yml b/.github/workflows/spend-logs-guard.yml new file mode 100644 index 00000000000..80a503679ec --- /dev/null +++ b/.github/workflows/spend-logs-guard.yml @@ -0,0 +1,49 @@ +name: Spend logs guard + +on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's diff is read as data and never executed + pull_request_target: + branches: + - main + pull_request_review: + types: + - submitted + - dismissed + +permissions: + contents: read + pull-requests: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + spend-logs-guard: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + ref: ${{ github.event.pull_request.base.ref }} + persist-credentials: false + - name: Fetch the pull request head, its merge base and its reviews + id: revisions + env: + GH_TOKEN: ${{ github.token }} + BASE_SHA: ${{ github.event.pull_request.base.sha }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + PR_NUMBER: ${{ github.event.pull_request.number }} + run: | + merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')" + git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA" + gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}/reviews" --paginate --slurp | jq 'add // []' > "$RUNNER_TEMP/reviews.json" + echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT" + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + - name: Run the guard + env: + MERGE_BASE: ${{ steps.revisions.outputs.merge_base }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + uv run --frozen python ci_cd/spend_logs_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --approver kerry-berri --reviews "$RUNNER_TEMP/reviews.json" diff --git a/ci_cd/spend_logs_guard.py b/ci_cd/spend_logs_guard.py new file mode 100644 index 00000000000..a04999c47c0 --- /dev/null +++ b/ci_cd/spend_logs_guard.py @@ -0,0 +1,119 @@ +"""Required check: a pull request that touches LiteLLM_SpendLogs in schema.prisma or a migration +needs an approving review from the spend-log owner before it can merge. + +CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies, +so a path rule would put every migration on the owner. This gate reads the diff instead: when no +added line in the watched files mentions the table it passes at once, otherwise it passes only once +the owner's latest review on the pull request is an approval. The workflow re-runs on every review. +""" + +from __future__ import annotations + +import argparse +import re +import subprocess +import sys +from collections.abc import Sequence +from pathlib import Path +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +WATCHED: Final = ( + "schema.prisma", + "litellm/proxy/schema.prisma", + "litellm-proxy-extras/litellm_proxy_extras/schema.prisma", +) +MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/" +TABLE: Final = re.compile(r"(? bool: + return path in WATCHED or (path.startswith(MIGRATIONS_PREFIX) and path.endswith("/migration.sql")) + + +def _added_lines(section: str) -> tuple[str, ...]: + body: Final = section.split("\n@@", 1) + return tuple(line[1:] for line in body[1].splitlines() if line.startswith("+")) if len(body) == 2 else () + + +def spend_log_hits(diff: str) -> tuple[str, ...]: + """`path: added line` for every added line in a watched file that names the table.""" + sections: Final = re.split(r"^diff --git [^\n]*\n", diff, flags=re.MULTILINE) + return tuple(f"{path}: {line.strip()}" for path, line in _watched_added_lines(sections) if TABLE.search(line)) + + +def _watched_added_lines(sections: Sequence[str]) -> tuple[tuple[str, str], ...]: + pairs: Final = tuple((_path_of(section), section) for section in sections) + return tuple( + (path, line) for path, section in pairs if watched(path) for line in _added_lines(section) + ) # comprehension-ok: path then its lines + + +def _path_of(section: str) -> str: + match: Final = re.search(r"^\+\+\+ b/(.+)$", section, flags=re.MULTILINE) + return match.group(1) if match else "" + + +def approved_by(reviews: Sequence[Review], login: str) -> bool: + """True when `login`'s most recent review on the pull request is an approval.""" + states: Final = tuple(review.state for review in reviews if review.user.login == login) + return bool(states) and states[-1] == "APPROVED" + + +def added_diff(base: str, head: str) -> str: + result: Final = subprocess.run( + ("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX), + check=True, + capture_output=True, + text=True, + ) + return result.stdout + + +class Arguments(BaseModel): + model_config = ConfigDict(frozen=True) + base: str + head: str + approver: str + reviews: Path + + +def main(argv: Sequence[str]) -> int: + parser: Final = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", required=True, help="merge base of the pull request") + parser.add_argument("--head", required=True, help="head sha of the pull request") + parser.add_argument("--approver", required=True, help="GitHub login whose approval unlocks the merge") + parser.add_argument("--reviews", required=True, help="JSON file with the pull request's reviews") + args: Final = Arguments.model_validate(vars(parser.parse_args(argv))) + hits: Final = spend_log_hits(added_diff(args.base, args.head)) + if not hits: + sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n") + return 0 + if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver): + sys.stdout.write(f"LiteLLM_SpendLogs change approved by {args.approver}: {', '.join(hits)}\n") + return 0 + listed: Final = "".join(f" {hit}\n" for hit in hits) + sys.stdout.write( + f"This pull request changes LiteLLM_SpendLogs and needs an approving review from @{args.approver}:\n" + f"{listed}The check re-runs when a review is submitted.\n" + ) + return 1 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/tests/unit/test_spend_logs_guard.py b/tests/unit/test_spend_logs_guard.py new file mode 100644 index 00000000000..9a033045dd0 --- /dev/null +++ b/tests/unit/test_spend_logs_guard.py @@ -0,0 +1,79 @@ +import importlib.util +import sys +from pathlib import Path +from typing import Final + +_SCRIPT: Final = Path(__file__).resolve().parents[2] / "ci_cd" / "spend_logs_guard.py" +_SPEC: Final = importlib.util.spec_from_file_location("spend_logs_guard", _SCRIPT) +assert _SPEC is not None and _SPEC.loader is not None +guard: Final = importlib.util.module_from_spec(_SPEC) +sys.modules["spend_logs_guard"] = guard +_SPEC.loader.exec_module(guard) + + +def _diff(path: str, added: str, start: int = 10) -> str: + return f"diff --git a/{path} b/{path}\n--- a/{path}\n+++ b/{path}\n@@ -{start},0 +{start},1 @@\n+{added}\n" + + +class TestSpendLogHits: + def test_an_added_schema_line_naming_the_table_is_a_hit(self): + diff: Final = _diff( + "schema.prisma", ' @@index([api_key, startTime], map: "LiteLLM_SpendLogs_api_key_startTime_idx")', 740 + ) + assert guard.spend_log_hits(diff) == ( + 'schema.prisma: @@index([api_key, startTime], map: "LiteLLM_SpendLogs_api_key_startTime_idx")', + ) + + def test_an_added_migration_line_is_a_hit_in_any_case(self): + path: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/20260101000000_x/migration.sql" + hits: Final = guard.spend_log_hits(_diff(path, 'CREATE INDEX "i" ON "litellm_spendlogs"("api_key");', 3)) + assert hits == (f'{path}: CREATE INDEX "i" ON "litellm_spendlogs"("api_key");',) + + def test_a_removed_line_or_another_table_or_an_unwatched_file_is_not_a_hit(self): + removed: Final = "--- a/schema.prisma\n+++ b/schema.prisma\n@@ -1,1 +1,0 @@\n-model LiteLLM_SpendLogs {\n" + other: Final = _diff("schema.prisma", " team_id String // LiteLLM_TeamTable, not MyLiteLLM_SpendLogs", 5) + elsewhere: Final = _diff("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5) + assert guard.spend_log_hits(removed + other + elsewhere) == () + + def test_two_files_in_one_diff_are_attributed_to_their_own_path(self): + diff: Final = _diff("litellm/proxy/schema.prisma", " a LiteLLM_SpendLogs", 1) + _diff( + "schema.prisma", " b String", 1 + ) + assert guard.spend_log_hits(diff) == ("litellm/proxy/schema.prisma: a LiteLLM_SpendLogs",) + + +class TestApprovedBy: + def test_the_latest_review_of_the_approver_decides(self): + reviews: Final = tuple( + guard.Review(state=state, user=guard.Reviewer(login=login)) + for login, state in ( + ("kerry-berri", "CHANGES_REQUESTED"), + ("someone-else", "APPROVED"), + ("kerry-berri", "APPROVED"), + ) + ) + assert guard.approved_by(reviews, "kerry-berri") is True + assert guard.approved_by(reviews[::-1], "kerry-berri") is False + + def test_no_review_from_the_approver_is_not_an_approval(self): + other: Final = guard.Review(state="APPROVED", user=guard.Reviewer(login="someone-else")) + assert guard.approved_by((other,), "kerry-berri") is False + assert guard.approved_by((), "kerry-berri") is False + + +class TestExitCode: + def test_a_pull_request_without_spend_log_lines_passes_without_reviews(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " foo String", 1)) + reviews: Final = tmp_path / "r.json" + reviews.write_text("[]") + assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 0 + assert "nothing to review" in capsys.readouterr().out + + def test_a_spend_log_change_fails_until_the_approver_approves(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " x LiteLLM_SpendLogs", 7)) + reviews: Final = tmp_path / "r.json" + reviews.write_text('[{"user": {"login": "someone-else"}, "state": "APPROVED"}]') + assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 + assert "schema.prisma: x LiteLLM_SpendLogs" in capsys.readouterr().out + reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED"}]') + assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 0 From 1bff4e3dae792f5884e9c366376ee147f94fd97e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 12:57:21 -0700 Subject: [PATCH 014/147] fix(anthropic): make opening usage estimates opt-in (#48633) * fix(anthropic): make opening usage estimates opt-in * docs(anthropic): follow README paragraph style --- litellm/__init__.py | 1 + litellm/anthropic_interface/readme.md | 36 ++++++++++- .../messages/adapters/streaming_iterator.py | 8 ++- .../test_streaming_iterator_message_id.py | 64 +++++++++++++++---- 4 files changed, 94 insertions(+), 15 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index c74292c5888..41e3fa34323 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -234,6 +234,7 @@ bedrock_neutralize_orphaned_tool_blocks: bool = True use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) ) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API +anthropic_messages_estimate_initial_input_tokens: bool = False # When True, strip the OpenAI-flavored `usage.total_tokens` field that # LiteLLM injects into non-streaming /v1/messages responses, bringing the # wire response into line with the Anthropic spec (matches the streaming diff --git a/litellm/anthropic_interface/readme.md b/litellm/anthropic_interface/readme.md index a864e2572e6..99e18e17de0 100644 --- a/litellm/anthropic_interface/readme.md +++ b/litellm/anthropic_interface/readme.md @@ -57,6 +57,40 @@ async for chunk in response: print(chunk) ``` +### Optional opening token estimate + +When `/v1/messages` streams through the Chat Completions adapter, the opening +`message_start.usage.input_tokens` defaults to `0`. To show an estimated prompt +size while a response streams, enable `anthropic_messages_estimate_initial_input_tokens`. +This setting defaults to `false` and applies to all requests using this adapter + +For an existing proxy deployment, add the setting under `litellm_settings` in +`config.yaml` and restart the proxy. Use a YAML boolean, without quotes + +```yaml +litellm_settings: + anthropic_messages_estimate_initial_input_tokens: true +``` + +For the Python SDK, set the option before making streaming Messages requests + +```python +import litellm + +litellm.anthropic_messages_estimate_initial_input_tokens = True +``` + +Send a streaming Messages request to a Chat Completions-backed model, such as a +`hosted_vllm/` model. The opening `message_start` event includes the estimated +input count. If estimation fails, the count remains `0`. Set the option to +`false` in YAML or `False` in Python to disable estimation and skip token counting + +The opening count is approximate and is not suitable for billing or enforcing +exact context limits. Opening output and cache counts remain `0`. Final +`message_delta` usage still comes from the provider and takes precedence over the +opening estimate. Native Anthropic streams, the Responses adapter, and +non-streaming responses are unaffected. This setting does not select the adapter + ### LiteLLM Proxy Server @@ -113,4 +147,4 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/messages' \ ], "max_tokens": 100 }' -``` \ No newline at end of file +``` diff --git a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index d1853bb6bd5..62a1e5f0a97 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py @@ -761,13 +761,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): def _create_initial_usage_delta(self) -> UsageDelta: """ - Estimate prompt tokens for message_start without consuming upstream chunks. + Optionally estimate prompt tokens without consuming upstream chunks. Output and cache counts start at zero. Provider usage in message_delta remains authoritative. Failed estimates fall back to zero. """ return UsageDelta( - input_tokens=self._estimate_initial_input_tokens(), + input_tokens=( + self._estimate_initial_input_tokens() + if litellm.anthropic_messages_estimate_initial_input_tokens is True + else 0 + ), output_tokens=0, cache_creation_input_tokens=0, cache_read_input_tokens=0, diff --git a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py index ec9146bc023..985cd38ef3c 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_message_id.py @@ -9,6 +9,7 @@ import datetime import json from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import Final +from unittest.mock import Mock import pytest import respx @@ -24,7 +25,6 @@ from litellm.llms.anthropic.messages.adapters.streaming_iterator import ( 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]) @@ -42,7 +42,9 @@ _BLOCK_PROMPTS: Final[tuple[list[AllMessageValues], ...]] = ( ) -def _assert_estimated_start_and_provider_final_usage(raw_events: Sequence[bytes], request_body: bytes) -> None: +def _assert_start_and_provider_final_usage( + raw_events: Sequence[bytes], request_body: bytes, estimate_enabled: bool +) -> 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( @@ -56,7 +58,7 @@ def _assert_estimated_start_and_provider_final_usage(raw_events: Sequence[bytes] ) assert expected > 0 assert start["usage"] == { - "input_tokens": expected, + "input_tokens": expected if estimate_enabled else 0, "output_tokens": 0, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, @@ -64,7 +66,12 @@ def _assert_estimated_start_and_provider_final_usage(raw_events: Sequence[bytes] 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: +@pytest.mark.parametrize("estimate_enabled", [None, False, True, "false"]) +def test_sync_message_start_opt_in_preserves_provider_final_usage( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, estimate_enabled: bool | str | None +) -> None: + if estimate_enabled is not None: + monkeypatch.setattr(litellm, "anthropic_messages_estimate_initial_input_tokens", estimate_enabled) stream: Final = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( max_tokens=16, messages=MESSAGES, @@ -74,13 +81,18 @@ def test_sync_message_start_estimates_prompt_without_replacing_provider_usage(re stream=True, custom_llm_provider="groq", ) - _assert_estimated_start_and_provider_final_usage(list(stream), respx_mock.calls.last.request.content) + _assert_start_and_provider_final_usage( + list(stream), respx_mock.calls.last.request.content, estimate_enabled is True + ) @pytest.mark.asyncio -async def test_async_message_start_estimates_prompt_without_replacing_provider_usage( - respx_mock: respx.MockRouter, +@pytest.mark.parametrize("estimate_enabled", [None, False, True, "false"]) +async def test_async_message_start_opt_in_preserves_provider_final_usage( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, estimate_enabled: bool | str | None ) -> None: + if estimate_enabled is not None: + monkeypatch.setattr(litellm, "anthropic_messages_estimate_initial_input_tokens", estimate_enabled) stream: Final = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( max_tokens=16, messages=MESSAGES, @@ -91,7 +103,7 @@ async def test_async_message_start_estimates_prompt_without_replacing_provider_u 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) + _assert_start_and_provider_final_usage(raw, respx_mock.calls.last.request.content, estimate_enabled is True) class _FailingUpstream: @@ -102,10 +114,31 @@ class _FailingUpstream: raise RuntimeError("upstream failed before its first chunk") +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_disabled_initial_estimate_skips_token_counting(monkeypatch: pytest.MonkeyPatch, is_async: bool) -> None: + counter: Final = Mock(side_effect=AssertionError("disabled estimation must not count tokens")) + monkeypatch.setattr(litellm, "anthropic_messages_estimate_initial_input_tokens", False) + monkeypatch.setattr(litellm, "token_counter", counter) + wrapper: Final = AnthropicStreamWrapper( + _FailingUpstream(), model="kimi-k2", completion_request={"messages": MESSAGES} + ) + first: Final = _EVENT.validate_python(await anext(wrapper) if is_async else 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, + } + counter.assert_not_called() + + @pytest.mark.parametrize("messages", _BLOCK_PROMPTS) def test_sync_initial_estimate_counts_content_blocks( - messages: list[AllMessageValues], respx_mock: respx.MockRouter + messages: list[AllMessageValues], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: + monkeypatch.setattr(litellm, "anthropic_messages_estimate_initial_input_tokens", True) 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) @@ -123,8 +156,9 @@ def test_sync_initial_estimate_counts_content_blocks( @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 + messages: list[AllMessageValues], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: + monkeypatch.setattr(litellm, "anthropic_messages_estimate_initial_input_tokens", True) 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) @@ -140,7 +174,10 @@ async def test_async_initial_estimate_counts_content_blocks( @pytest.mark.parametrize("prompt", [{}, {"messages": []}, {"messages": "invalid"}]) -def test_sync_initial_estimate_fallback_does_not_consume_upstream(prompt: Mapping[str, object]) -> None: +def test_sync_initial_estimate_fallback_does_not_consume_upstream( + prompt: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "anthropic_messages_estimate_initial_input_tokens", True) wrapper: Final = AnthropicStreamWrapper(_FailingUpstream(), model="kimi-k2", completion_request=prompt) first: Final = _EVENT.validate_python(next(wrapper)) message: Final = _EVENT.validate_python(first["message"]) @@ -155,7 +192,10 @@ def test_sync_initial_estimate_fallback_does_not_consume_upstream(prompt: Mappin @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: +async def test_async_initial_estimate_fallback_does_not_consume_upstream( + prompt: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "anthropic_messages_estimate_initial_input_tokens", True) 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"]) From 4846d9c2ed38795b931cd397e70705e45090b280 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 20:24:56 +0000 Subject: [PATCH 015/147] fix(cost-map): add 2026-11-06 deprecation_date to two together_ai rows (#48651) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7cd4eadca6e..39500330096 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -47102,6 +47102,7 @@ }, "together_ai/MiniMaxAI/MiniMax-M3": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-11-06", "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, @@ -47295,6 +47296,7 @@ }, "together_ai/thinkingmachines/Inkling": { "cache_read_input_token_cost": 1.7e-07, + "deprecation_date": "2026-11-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 524288, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7cd4eadca6e..39500330096 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -47102,6 +47102,7 @@ }, "together_ai/MiniMaxAI/MiniMax-M3": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-11-06", "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, @@ -47295,6 +47296,7 @@ }, "together_ai/thinkingmachines/Inkling": { "cache_read_input_token_cost": 1.7e-07, + "deprecation_date": "2026-11-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 524288, From d6edd37810944752561ed49d9a8e85085f8dce81 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 13:33:03 -0700 Subject: [PATCH 016/147] fix(aiohttp_transport): close an owned aiohttp session when its transport is collected (#48650) An OpenAI or Azure SDK client litellm built and evicted from the client cache is usually garbage collected inside the closer's 900 second grace window, so the closer's weak reference dies and the client's aiohttp session is never closed. LiteLLMAiohttpTransport.__del__ now hands an owned, still-open ClientSession to _close_recycled_session, which already covers the current-loop, foreign-loop and closed-loop cases. Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: BangerSoul --- .../llms/custom_httpx/aiohttp_transport.py | 7 ++ .../sdk/test_aiohttp_session_rebuild_wire.py | 86 ++++++++++++++++++- .../unit/caching/test_llm_caching_handler.py | 33 +++++++ .../custom_httpx/test_aiohttp_transport.py | 73 ++++++++++++++++ 4 files changed, 197 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b8cb183c02f..0e4d97fd770 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -289,6 +289,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport): cls._background_close_tasks.add(task) task.add_done_callback(cls._on_close_task_done) + def __del__(self) -> None: + session: Final = getattr(self, "client", None) + if not getattr(self, "_owns_session", False) or not isinstance(session, ClientSession): + return + with contextlib.suppress(Exception): + self._close_recycled_session(session) + def get_valid_client_session(self) -> ClientSession: """ Helper to get a valid ClientSession for the current event loop. diff --git a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py index 1ec1f486bed..f1a18ba389a 100644 --- a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py +++ b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py @@ -11,11 +11,30 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final import pytest -from pydantic import JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter CONFIGURED_KEEPALIVE_SECONDS: Final = 1 IDLE_SECONDS: Final = 2 +EVICTED_CLIENTS: Final = 30 RESPONSES: Final = TypeAdapter(list[dict[str, JsonValue]]) +CHAT_COMPLETION: Final = { + "id": "chatcmpl-scripted", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class EvictionReport(BaseModel): + model_config = ConfigDict(frozen=True) + + cache_size: int + cached_clients: int + established_connections: int + unclosed_messages: tuple[str, ...] + REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( """ @@ -38,6 +57,47 @@ REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( ) +EVICTED_CLIENTS_EXCHANGE: Final = textwrap.dedent( + """ + import asyncio, gc, json, sys + import psutil + import litellm + from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport + + unclosed = [] + + def established(port: int) -> int: + connections = psutil.Process().net_connections("tcp") + return sum(1 for c in connections if c.status == "ESTABLISHED" and c.raddr and c.raddr.port == port) + + async def one_call(base_url: str, index: int) -> None: + await litellm.acompletion( + model="openai/gpt-5.4-nano", + api_key=f"sk-distinct-{index}", + api_base=base_url, + messages=[{"role": "user", "content": "hi"}], + ) + + async def main(base_url: str, port: int, evicted: int) -> None: + asyncio.get_running_loop().set_exception_handler(lambda loop, ctx: unclosed.append(ctx["message"])) + cache = litellm.in_memory_llm_clients_cache + for index in range(cache.max_size_in_memory + evicted): + await one_call(base_url, index) + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + await asyncio.sleep(0) + print(json.dumps({ + "cache_size": cache.max_size_in_memory, + "cached_clients": len(cache.cache_dict), + "established_connections": established(port), + "unclosed_messages": unclosed, + })) + + asyncio.run(main(sys.argv[1], int(sys.argv[2]), int(sys.argv[3]))) + """ +) + + class _ConnectionCountingPeer(ThreadingHTTPServer): daemon_threads = True @@ -62,7 +122,8 @@ class _ConnectionHandler(BaseHTTPRequestHandler): def do_POST(self) -> None: self.rfile.read(int(self.headers["Content-Length"])) - body: Final = json.dumps({"connection": self.connection_number}).encode() + payload: Final = CHAT_COMPLETION if self.path.endswith("/chat/completions") else {"connection": self.connection_number} + body: Final = json.dumps(payload).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) @@ -107,3 +168,24 @@ def test_rebuilt_shared_session_drops_idle_connection_after_configured_keepalive ) -> None: observed: Final = _rebuilt_session_exchange(connection_counting_peer) assert observed == [{"connection": 1}, {"connection": 2}], observed + + +def _evicted_clients_exchange(base_url: str) -> EvictionReport: + completed: Final = subprocess.run( + [sys.executable, "-P", "-c", EVICTED_CLIENTS_EXCHANGE, base_url, base_url.rsplit(":", 1)[1], str(EVICTED_CLIENTS)], + capture_output=True, + text=True, + timeout=300, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return EvictionReport.model_validate_json(completed.stdout) + + +def test_sdk_clients_evicted_from_a_full_client_cache_leave_no_open_session_or_connection( + connection_counting_peer: str, +) -> None: + report: Final = _evicted_clients_exchange(connection_counting_peer) + assert report.cached_clients == report.cache_size, report + assert report.unclosed_messages == (), report + assert report.established_connections <= report.cache_size, report diff --git a/tests/unit/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py index 0c0fe8ca3c5..a3efd5d97aa 100644 --- a/tests/unit/caching/test_llm_caching_handler.py +++ b/tests/unit/caching/test_llm_caching_handler.py @@ -11,11 +11,14 @@ See: https://github.com/BerriAI/litellm/pull/22247 import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref import warnings +import aiohttp import pytest +from openai import AsyncOpenAI from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import( @@ -204,6 +207,36 @@ async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elap assert client.closed is True +@pytest.mark.asyncio +async def test_an_evicted_owned_client_collected_inside_its_grace_window_still_closes_its_aiohttp_session(): + """ + The closer holds an evicted client weakly for the whole grace window. An SDK + process that exits, or a proxy that evicts and never calls the cache again, + collects the client first, so only its own finalizer can still close the + aiohttp session (LIT-9767). + """ + clock = _FakeClock() + cache = LLMClientCache( + max_size_in_memory=2, + evicted_client_closer=EvictedClientCloser(grace_seconds=900.0, clock=clock), + ) + session = aiohttp.ClientSession() + client = AsyncOpenAI( + api_key="sk-evicted", + http_client=httpx.AsyncClient(transport=LiteLLMAiohttpTransport(client=session)), + ) + cache.set_cache("client-key", client, litellm_owned_client=True, ttl=600) + + cache.ttl_dict = {key: 0 for key in cache.ttl_dict} + cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap] + cache.evict_cache() + del client + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the evicted client was collected inside the grace window and left its session open" + + @pytest.mark.asyncio async def test_evicted_caller_supplied_client_is_never_closed(): """litellm does not own a client the caller passed in, so it must stay open.""" diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index 0765371dc4a..c7d5518ef5b 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -1,8 +1,10 @@ import aiohttp as aiohttp_aiohttp_handler, asyncio, importlib, litellm import concurrent.futures +import gc import socket import sys import struct +import threading from typing import Final import aiohttp @@ -1310,3 +1312,74 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat assert len(received_requests) == 1 finally: await handler.close() + + +async def _new_session() -> aiohttp.ClientSession: + return aiohttp.ClientSession() + + +async def _until_closed(session: aiohttp.ClientSession) -> None: + while not session.closed: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_collected_transport_closes_the_session_it_owns(): + """An SDK client evicted from the cache is collected long before the closer's grace window ends.""" + session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the owned session outlived the transport that was collected before closing it" + + +@pytest.mark.asyncio +async def test_a_collected_transport_leaves_a_shared_session_open(): + session = aiohttp.ClientSession() + try: + transport = LiteLLMAiohttpTransport(client=session, owns_session=False) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert not session.closed, "a transport closed a session it does not own" + finally: + await session.close() + + +def test_a_transport_collected_after_its_loop_closed_still_closes_its_session(caplog: pytest.LogCaptureFixture): + """asyncio.run() closes the loop the session lives on before the evicted client is collected.""" + loop = asyncio.new_event_loop() + session = loop.run_until_complete(_new_session()) + loop.close() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + + assert session.closed, "the owned session stayed open after its transport was collected on a closed loop" + del session + gc.collect() + assert "Unclosed client session" not in caplog.text, caplog.text + + +def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_on_that_loop(): + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + session = asyncio.run_coroutine_threadsafe(_new_session(), loop).result(timeout=10) + transport = LiteLLMAiohttpTransport(client=session) + try: + del transport + gc.collect() + asyncio.run_coroutine_threadsafe(_until_closed(session), loop).result(timeout=10) + + assert session.closed, "the owned session stayed open after its transport was collected from another thread" + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=10) + loop.close() From 9e01b3c37b0736806aa70dcde72f5d85079be3e8 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 13:33:32 -0700 Subject: [PATCH 017/147] fix(complexity_router): keep the held user_turn route for Codex compaction requests (#48627) * test(complexity_router): Codex compaction requests hold the user_turn decision Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(complexity_router): treat Codex compaction traffic as user_turn continuation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(complexity_router): validate Codex turn metadata with frozen models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(complexity_router): flatten body system text through _message_text Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(complexity_router): drop unrelated reformatting from the mapped test file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(complexity_router): skip pin escalation on Codex compaction turns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../complexity_router/complexity_router.py | 97 ++++++++++++++-- ...test_complexity_router_codex_compaction.py | 109 ++++++++++++++++++ .../router_strategy/test_complexity_router.py | 96 +++++++++++++++ 3 files changed, 290 insertions(+), 12 deletions(-) create mode 100644 tests/integration/routing/test_complexity_router_codex_compaction.py diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index cc592bbcbc5..12968bf7d4a 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -28,7 +28,7 @@ from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast -from pydantic import BaseModel, TypeAdapter, ValidationError, create_model +from pydantic import BaseModel, Field, TypeAdapter, ValidationError, create_model from pydantic_core import ErrorDetails from typing_extensions import ReadOnly, TypedDict @@ -773,6 +773,56 @@ def _newest_turn_is_human_ask( return _last_human_ask_index(messages, marker_pairs) == newest_activity +_CODEX_TURN_METADATA_KEY: Final = "x-codex-turn-metadata" +_CODEX_COMPACTION_SUMMARY_PREFIX: Final = ( + "Another language model started to solve this problem and produced a summary of its thinking process." +) + + +class _CodexTurnMetadata(BaseModel, frozen=True): + request_kind: str | None = None + + +class _CodexTurnMetadataCarrier(BaseModel, frozen=True): + turn_metadata: str | None = Field(default=None, alias=_CODEX_TURN_METADATA_KEY) + + +class _CodexProxyRequest(BaseModel, frozen=True): + headers: _CodexTurnMetadataCarrier | None = None + + +class _CodexRequestMetadata(BaseModel, frozen=True): + """Codex sends its turn metadata in the Responses body's `client_metadata` and projects it onto a + header of the same name for transports without that field""" + + client_metadata: _CodexTurnMetadataCarrier | None = None + proxy_server_request: _CodexProxyRequest | None = None + + def turn_metadata_values(self) -> tuple[str | None, ...]: + headers: Final = self.proxy_server_request.headers if self.proxy_server_request else None + return ( + self.client_metadata.turn_metadata if self.client_metadata else None, + headers.turn_metadata if headers else None, + ) + + +def _codex_request_kind(value: str | None) -> str | None: + if value is None: + return None + try: + return _CodexTurnMetadata.model_validate_json(value).request_kind + except ValidationError: + return None + + +def _is_codex_compaction_request(request_kwargs: Mapping[str, object]) -> bool: + try: + request_metadata: Final = _CodexRequestMetadata.model_validate(request_kwargs) + except ValidationError: + return False + return any(_codex_request_kind(value) == "compaction" for value in request_metadata.turn_metadata_values()) + + def _iter_system_scope_texts( body_system: object, messages: Sequence[Mapping[str, object]], @@ -787,10 +837,9 @@ def _iter_system_scope_texts( must go through the staleness-aware tail scan instead -- scanning it here would floor every turn of a session that once planned, for any pattern whose client injects mid-conversation. """ - if isinstance(body_system, str): - yield body_system - elif isinstance(body_system, list): - yield _message_text(body_system) + body_system_text: Final = _message_text(body_system) + if body_system_text: + yield body_system_text for msg in messages: if msg.get("role") != "system": return @@ -4263,14 +4312,31 @@ class ComplexityRouter(CustomLogger): def _reminder_markers_for_request(self, request_kwargs: Mapping[str, object]) -> tuple[tuple[str, str], ...]: if self.config.reminder_markers is not None: return self._reminder_markers - if any( + return _CODEX_REMINDER_MARKERS if self._is_codex_request(request_kwargs) else _DEFAULT_REMINDER_MARKERS + + @staticmethod + def _is_codex_request(request_kwargs: Mapping[str, object]) -> bool: + return any( is_codex_user_agent(user_agent) for metadata_key in ("litellm_metadata", "metadata") if isinstance(metadata := request_kwargs.get(metadata_key), Mapping) if isinstance(user_agent := metadata.get("user_agent"), str) - ): - return _CODEX_REMINDER_MARKERS - return _DEFAULT_REMINDER_MARKERS + ) + + def _is_codex_compaction_lifecycle( + self, + request_kwargs: Mapping[str, object], + resolved_messages: Sequence[Mapping[str, object]] | None, + marker_pairs: tuple[tuple[str, str], ...], + ) -> bool: + """Codex's compaction call and the first request after it carry a synthetic user turn (the + summarization prompt, then the summary), so `user_turn` must read both as continuations.""" + if _is_codex_compaction_request(request_kwargs): + return True + if not resolved_messages or not self._is_codex_request(request_kwargs): + return False + newest_ask: Final = _newest_turn_ask(resolved_messages, marker_pairs) + return newest_ask is not None and newest_ask.startswith(_CODEX_COMPACTION_SUMMARY_PREFIX) @staticmethod def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: @@ -4380,15 +4446,22 @@ class ComplexityRouter(CustomLogger): # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human # ask falls through and re-classifies. session_affinity restores pin-first for asks too. - pin_replay_allowed: Final = bool(self.config.session_affinity) or not _newest_turn_is_human_ask( - resolved_messages, marker_pairs + codex_compaction: Final = self._is_codex_compaction_lifecycle(request_kwargs, resolved_messages, marker_pairs) + pin_replay_allowed: Final = ( + bool(self.config.session_affinity) + or not _newest_turn_is_human_ask(resolved_messages, marker_pairs) + or codex_compaction ) if cache_key is not None and pin_replay_allowed: pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: - user_message: Final = _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None + user_message: Final = ( + _newest_turn_ask(resolved_messages, marker_pairs) + if resolved_messages and not codex_compaction + else None + ) pin_escalation_keyword: Final = ( self._matched_escalation_keyword(user_message) if user_message is not None else None ) diff --git a/tests/integration/routing/test_complexity_router_codex_compaction.py b/tests/integration/routing/test_complexity_router_codex_compaction.py new file mode 100644 index 00000000000..2f436901cd0 --- /dev/null +++ b/tests/integration/routing/test_complexity_router_codex_compaction.py @@ -0,0 +1,109 @@ +import json +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +_ASK: Final = {"role": "user", "content": "Design a Byzantine fault tolerant consensus protocol and prove it safe"} +_COMPACTION_PROMPT: Final = {"role": "user", "content": "Hello! Summarize the conversation so far."} + + +def _completion(content: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex[:8], + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, + } + ).encode() + ) + + +def _answering(content: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + return _completion(content) + + return respond + + +def _classifier(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + body: Final = request.body.decode() + tier: Final = "SIMPLE" if "Summarize" in body else "REASONING" + return _completion(json.dumps({"tier": tier})) + + +def test_codex_compaction_request_stays_on_the_held_user_turn_model(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(_classifier) as judge, + wire_server(_answering("simple answer")) as simple, + wire_server(_answering("reasoning answer")) as reasoning, + ): + router: Final = "router-" + uuid.uuid4().hex[:8] + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + { + "model_name": name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": url, "api_key": "synthetic"}, + } + for name, url in (("judge", judge.url), ("simple", simple.url), ("reasoning", reasoning.url)) + ] + [ + { + "model_name": router, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "judge", "timeout_ms": 20000}, + "classification_mode": "user_turn", + "tiers": { + "SIMPLE": "simple", + "MEDIUM": "simple", + "COMPLEX": "reasoning", + "REASONING": "reasoning", + }, + }, + }, + } + ] + path: Final = tmp_path / "complexity_router.yaml" + path.write_text(yaml.safe_dump(config)) + session: Final = uuid.uuid4().hex + + def codex_headers(request_kind: str) -> dict[str, str]: + return { + "User-Agent": "codex-tui/0.164.0", + "session_id": session, + "x-codex-turn-metadata": json.dumps({"session_id": session, "request_kind": request_kind}), + } + + with owned_proxy(gateway, tmp_path, {}, config=path) as candidate: + ask: Final = candidate.request( + "POST", "/v1/chat/completions", {"model": router, "messages": [_ASK]}, headers=codex_headers("turn") + ) + assert ask.status_code == 200, ask.text + assert ask.json()["choices"][0]["message"]["content"] == "reasoning answer", ask.text + compaction: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": router, "messages": [_ASK, {"role": "assistant", "content": "done"}, _COMPACTION_PROMPT]}, + headers=codex_headers("compaction"), + ) + assert compaction.status_code == 200, compaction.text + assert compaction.json()["choices"][0]["message"]["content"] == "reasoning answer", compaction.text + assert len([call for call in judge.drain() if call.method == "POST"]) == 1 + assert [call for call in simple.drain() if call.method == "POST"] == [] diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index a46f415e10a..2d27a025ffa 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -7820,6 +7820,102 @@ class TestClassificationMode: assert pure.routing_decision["cause"] == "user_turn_continuation" assert hybrid.model == "gpt-4o-mini" + CODEX_USER_AGENT = "codex-tui/0.164.0 (Ubuntu 22.4.0; x86_64) screen (codex-tui; 0.164.0)" + COMPACTION_PROMPT = {"role": "user", "content": "Hello! Summarize the conversation so far."} + COMPACTION_SUMMARY = { + "role": "user", + "content": "Another language model started to solve this problem and produced a summary of its thinking " + "process. You also have access to the state of the tools that were used by that language model.\nHi.", + } + + def _codex_kwargs( + self, session_id: str, request_kind: str, transport: str = "client_metadata" + ) -> dict[str, object]: + turn_metadata: Final = json.dumps({"session_id": session_id, "request_kind": request_kind}) + carrier: Final = ( + {"client_metadata": {"x-codex-turn-metadata": turn_metadata}} + if transport == "client_metadata" + else {"proxy_server_request": {"headers": {"x-codex-turn-metadata": turn_metadata}}} + ) + return {"metadata": {"session_id": session_id, "user_agent": self.CODEX_USER_AGENT}, **carrier} + + @pytest.mark.asyncio + @pytest.mark.parametrize("transport", ["client_metadata", "header"]) + async def test_codex_compaction_request_replays_held_decision( + self, mock_router_instance, user_turn_config, transport + ): + """Codex's compaction call carries a synthetic user prompt; it must hold the session's + decision instead of classifying the prompt as a new human ask.""" + router = self._router(mock_router_instance, user_turn_config) + session = f"codex-compact-{transport}" + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._codex_kwargs(session, "turn"), messages=[self.REASONING_ASK] + ) + compaction = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs(session, "compaction", transport), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.COMPACTION_PROMPT], + ) + assert compaction.model == "o1-preview" + assert compaction.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_codex_post_compaction_summary_is_a_continuation(self, mock_router_instance, user_turn_config): + """The first request after a mid-turn compaction ends on Codex's summary fragment, which + is not a human ask, so the held decision carries on.""" + router = self._router(mock_router_instance, user_turn_config) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-summary", "turn"), + messages=[self.REASONING_ASK], + ) + resumed = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-summary", "turn"), + messages=[self.REASONING_ASK, self.COMPACTION_SUMMARY], + ) + assert resumed.model == "o1-preview" + assert resumed.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_codex_summary_quoting_escalation_keyword_keeps_held_tier( + self, mock_router_instance, user_turn_config + ): + router = self._router(mock_router_instance, user_turn_config) + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._codex_kwargs("codex-quote", "turn"), messages=[self.SIMPLE_ASK] + ) + quoted_summary: Final = { + "role": "user", + "content": f"{self.COMPACTION_SUMMARY['content']} The user said LITELLM ESCALATE", + } + resumed = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-quote", "turn"), + messages=[self.SIMPLE_ASK, quoted_summary], + ) + assert resumed.model == "gpt-4o-mini" + assert resumed.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_codex_human_ask_after_compaction_reclassifies(self, mock_router_instance, user_turn_config): + router = self._router(mock_router_instance, user_turn_config) + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._codex_kwargs("codex-next", "turn"), messages=[self.REASONING_ASK] + ) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-next", "compaction"), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.COMPACTION_PROMPT], + ) + next_ask = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-next", "turn"), + messages=[self.REASONING_ASK, self.COMPACTION_SUMMARY, self.SIMPLE_ASK], + ) + assert next_ask.model == "gpt-4o-mini" + assert next_ask.routing_decision["cause"] != "user_turn_continuation" + @pytest.mark.asyncio async def test_session_affinity_wins_when_both_knobs_are_on(self, mock_router_instance, user_turn_config): """With session_affinity also on, the pin short-circuits new asks too and keeps its own From 34c10829b1f2aba5ee44887a413bac12d77e2bc0 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 13:33:47 -0700 Subject: [PATCH 018/147] fix(ui): percent-encode request id in spend log detail URL (#48630) * test: cover request ids with URL special characters in log detail fetches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): percent-encode request id in spend log detail URL Co-authored-by: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> --- litellm/proxy/client/cli/commands/debug.py | 3 ++- .../proxy/client/cli/test_debug_commands.py | 18 ++++++++++++++++ .../src/components/networking.test.ts | 21 +++++++++++++++++++ .../src/components/networking.tsx | 4 ++-- 4 files changed, 43 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/client/cli/commands/debug.py b/litellm/proxy/client/cli/commands/debug.py index 8728a9323b7..89b93e2bbb8 100644 --- a/litellm/proxy/client/cli/commands/debug.py +++ b/litellm/proxy/client/cli/commands/debug.py @@ -15,6 +15,7 @@ from datetime import datetime, timezone from pathlib import Path from types import MappingProxyType from typing import Final +from urllib.parse import quote import click import requests @@ -177,7 +178,7 @@ class SpendLogsFetcher: return DebugFailure(f"Unexpected {uri} response: {e}") def payload(self, request_id: str) -> RequestResponsePayload | None | DebugFailure: - uri: Final = f"/spend/logs/ui/{request_id}" + uri: Final = f"/spend/logs/ui/{quote(request_id, safe='')}" raw: Final = self._get(uri) if isinstance(raw, DebugFailure): return raw diff --git a/tests/unit/proxy/client/cli/test_debug_commands.py b/tests/unit/proxy/client/cli/test_debug_commands.py index 419fd821cfb..311edde0af9 100644 --- a/tests/unit/proxy/client/cli/test_debug_commands.py +++ b/tests/unit/proxy/client/cli/test_debug_commands.py @@ -1,6 +1,8 @@ import json import os +import re import time +from urllib.parse import unquote import pytest import requests @@ -114,6 +116,22 @@ def test_recent_bodies_fetches_latest_turns_even_when_successful(): assert _called_paths() == ["/spend/logs/session/ui", "/spend/logs/ui/req-ok"] +@pytest.mark.parametrize("request_id", ["req#a1", "req?a=1", "req%41b", "req a&b"]) +@responses.activate +def test_bodies_are_fetched_for_request_ids_with_url_special_characters(request_id): + def payload_for_stored_id(request): + stored_id = unquote(request.path_url.split("?")[0].removeprefix("/spend/logs/ui/")) + found = stored_id == request_id + return 200, {}, json.dumps({"response": {"id": "msg_special"}} if found else None) + + _mock_proxy([{**OK_ROW, "request_id": request_id}], {}) + responses.add_callback(responses.GET, re.compile(f"{PROXY}/spend/logs/ui/.*"), callback=payload_for_stored_id) + result = CliRunner().invoke(cli, ["debug", "claude", "--session-id", SESSION, "--no-save"]) + + assert result.exit_code == 0, result.output + assert "msg_special" in result.output + + @responses.activate def test_bodies_are_truncated_to_max_chars(): _mock_proxy([OK_ROW], {"req-ok": {"proxy_server_request": {"body": "a" * 5000}, "response": None}}) diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 8b0f0389772..874cdb52fad 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -602,6 +602,27 @@ describe("uiSpendLogsCall exclude_internal_health_checks serialization", () => { }); }); +describe("uiSpendLogDetailsCall", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it.each(["chatcmpl-123#a1", "chatcmpl-123?a=1", "chatcmpl-123%41b", "chatcmpl-123 a&b"])( + "requests the details of %s with the full request id and start date", + async (logId) => { + const mockFetch = vi.fn().mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue({ id: logId }) }); + vi.stubGlobal("fetch", mockFetch); + + await Networking.uiSpendLogDetailsCall("token", logId, "2026-01-01 00:00:00"); + + const url = new URL(String(mockFetch.mock.calls.at(-1)?.[0]), "http://example.com"); + expect(decodeURIComponent(url.pathname.split("/spend/logs/ui/")[1])).toBe(logId); + expect(url.hash).toBe(""); + expect([...url.searchParams.entries()]).toEqual([["start_date", "2026-01-01 00:00:00"]]); + }, + ); +}); + describe("sessionSpendLogsCall", () => { const originalFetch = global.fetch; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 38d150f1d63..486ec2be068 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4709,8 +4709,8 @@ export const uiSpendLogDetailsCall = async (accessToken: string, logId: string, try { // Construct base URL let url = proxyBaseUrl - ? `${proxyBaseUrl}/spend/logs/ui/${logId}?start_date=${encodeURIComponent(start_date)}` - : `/spend/logs/ui/${logId}?start_date=${encodeURIComponent(start_date)}`; + ? `${proxyBaseUrl}/spend/logs/ui/${encodeURIComponent(logId)}?start_date=${encodeURIComponent(start_date)}` + : `/spend/logs/ui/${encodeURIComponent(logId)}?start_date=${encodeURIComponent(start_date)}`; const response = await fetch(url, { method: "GET", From 52562c47f83be9df28ecf6ca9f9fb620286d04b9 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 13:53:23 -0700 Subject: [PATCH 019/147] refactor(proxy-extras): drop the spend-log api_key special case from the row check remedy (#48636) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/migration_data_checks.py | 8 -------- .../database/migrations/test_request_log_indexes.py | 9 ++++++--- .../litellm_proxy_extras/test_migration_data_checks.py | 2 +- 3 files changed, 7 insertions(+), 12 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py index a2d3728349a..7b01e64dd73 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py @@ -42,12 +42,6 @@ _OCTET_LENGTH_TYPES: Final = frozenset({"text", "varchar", "bpchar", "bytea", "n 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"] @@ -571,8 +565,6 @@ def find_blockers( 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": diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index bf37ca80665..1de2090b487 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -956,8 +956,11 @@ def _legacy_bearer_token() -> str: 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}$';")] +_HASH_LEGACY_KEYS: Final = ( + 'UPDATE "LiteLLM_SpendLogs" SET api_key = encode(sha256(convert_to(' + "regexp_replace(api_key, '^[Bb][Ee][Aa][Rr][Ee][Rr] ', ''), 'UTF8')), 'hex') " + "WHERE api_key <> '' AND api_key !~ '^[0-9a-fA-F]{64}$' AND api_key !~ '^hashed-jwt-[0-9A-Fa-f]{64}$'" +) @requires_db @@ -985,7 +988,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key 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 + conn.execute(_HASH_LEGACY_KEYS) 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) diff --git a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py index 7897a388474..f3b787beda1 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py +++ b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py @@ -125,7 +125,7 @@ def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_wa "(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 'fix: Shorten or hash the long "LiteLLM_SpendLogs" (api_key, startTime) values' in text assert text.startswith("Existing rows would make a pending schema change fail:") assert "No schema change was made" in text From e0d1e64591dc0f4243a17014ce216943ffa0d186 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:04:30 -0400 Subject: [PATCH 020/147] feat(router): respect caller output budgets in complexity routing (#48632) Co-authored-by: stuti612 --- litellm/router.py | 96 ++++---- .../complexity_router/README.md | 14 ++ .../complexity_router/complexity_router.py | 63 ++++-- .../complexity_router/config.py | 5 +- .../common_utils/test_cache_aware_routing.py | 18 +- .../router_strategy/test_complexity_router.py | 206 +++++++++++++++++- .../add_model/ComplexityRouterConfig.tsx | 2 +- .../add_model/ResponseFormatControls.tsx | 56 ++++- .../build_complexity_router_config.test.ts | 20 ++ .../build_complexity_router_config.ts | 8 +- ...d_updated_complexity_router_config.test.ts | 23 ++ ...dit_auto_router_modal.integration.test.tsx | 22 +- .../hydrate_complexity_router_config.ts | 3 +- .../src/lib/autorouter_presets.test.ts | 7 + .../src/lib/autorouter_presets.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 16 files changed, 448 insertions(+), 100 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 0d4045e6ca8..c8153b596d6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1122,6 +1122,27 @@ def as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None +def stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: + """None means unstamped; an empty mapping means the original caller supplied no limit.""" + stamped: Final = next( + ( + bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket + ), + None, + ) + if not isinstance(stamped, dict): + return None + return MappingProxyType( + { + carrier: cap + for carrier, value in stamped.items() + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None + } + ) + + def _exact_base_model_max_input_tokens(base_model: str) -> int | None: try: model_info: Final = litellm.get_model_info(model=base_model) @@ -15235,9 +15256,7 @@ class Router: request_kwargs: dict, responses_call: bool, ) -> bool: - """Apply a routing strategy's per-tier litellm_params on top of the request and report - whether they pinned an output ceiling, so the caller can hand the request its own ceiling - back on a routing pass that pins none.""" + """Apply per-tier litellm_params and report whether they pinned an output ceiling.""" if not tier_litellm_params: return False reject_api_key_secret_reference(tier_litellm_params) @@ -15250,34 +15269,26 @@ class Router: return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params) @staticmethod - def _restore_client_ceiling_no_tier_pins(request_kwargs: MutableMapping[str, object]) -> None: - """A model-group fallback re-enters routing with the kwargs an earlier auto-router pass - already rewrote, so a ceiling sized for that pass's tier would ride onto a group no tier - chose. When this pass pins none, hand the request back exactly the carriers the caller - sent, which the first pinning pass stamped. The stamp lives in a metadata bucket a - caller can also write, so the proxy strips the key at ingestion and this read takes - nothing but the three ceiling carriers as integers: no other key ever reaches kwargs.""" - stamped: Final = next( - ( - bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) - for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) - if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket - ), - None, - ) - if not isinstance(stamped, dict): + def _restore_client_output_ceiling(request_kwargs: MutableMapping[str, object]) -> None: + """Restore caller intent before routing, so neither a previous tier pin nor alias defaults + become the next attempt's budget. The proxy strips caller-written stamps at ingestion.""" + callers_ceiling: Final = stamped_client_ceiling(request_kwargs) + if callers_ceiling is None: return - callers_ceiling: Final = MappingProxyType( - { - carrier: cap - for carrier, value in stamped.items() - if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None - } - ) for carrier in OUTPUT_TOKEN_CEILING_PARAMS: request_kwargs.pop(carrier, None) request_kwargs.update(callers_ceiling) + @staticmethod + def _stamp_client_output_ceiling( + request_kwargs: dict[str, object], # mutable-ok: preserve request-owned metadata across routing attempts + ) -> None: + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket.setdefault( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + {carrier: request_kwargs[carrier] for carrier in OUTPUT_TOKEN_CEILING_PARAMS if carrier in request_kwargs}, + ) + @staticmethod def _drop_client_carriers_a_tier_pin_supersedes( request_kwargs: dict[str, object], @@ -15294,15 +15305,7 @@ class Router: one setting under three names, and a provider handed two of them either rejects the request or picks one by iteration order.""" if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): - _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket.setdefault( - CLIENT_OUTPUT_CEILING_METADATA_KEY, - { - carrier: request_kwargs[carrier] - for carrier in OUTPUT_TOKEN_CEILING_PARAMS - if carrier in request_kwargs - }, - ) + Router._stamp_client_output_ceiling(request_kwargs) for carrier in OUTPUT_TOKEN_CEILING_PARAMS: request_kwargs.pop(carrier, None) if "reasoning_effort" not in tier_litellm_params: @@ -15374,14 +15377,12 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + self._pin_tier_params_onto_request( model=model, tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, request_kwargs=request_kwargs, responses_call=responses_call, ) - if not tier_pins_ceiling: - self._restore_client_ceiling_no_tier_pins(request_kwargs) ######################################################### # Resolve the strategy and logger AFTER the pre-routing hook, since @@ -15539,14 +15540,12 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + self._pin_tier_params_onto_request( model=model, tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, request_kwargs=request_kwargs, responses_call=responses_call, ) - if not tier_pins_ceiling: - self._restore_client_ceiling_no_tier_pins(request_kwargs) # 2. Get healthy deployments healthy_deployments: Final = await self.async_get_healthy_deployments( @@ -15896,6 +15895,7 @@ class Router: spend metadata is stamped before routing and the response carries the tier group the strategy picked. """ + self._restore_client_output_ceiling(request_kwargs) self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) requested_registered_model_name: Final = self.get_model_from_alias(model=model) or model registered_model_name: Final = await self._resolve_claude_code_session_router( @@ -16062,7 +16062,8 @@ class Router: # them) until `_update_kwargs_with_deployment` drops any the selected # deployment sets itself (its own `aws_region_name` beats the marker's). # Per-tier `litellm_params` on the hook response are deliberate overrides - # the caller applies on top, so those keys are never forwarded here. + # the caller applies on top, so those keys and equivalent output-limit + # names are never tracked as alias defaults here. marker_params: Final = ( self._forwardable_alias_marker_params( model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs @@ -16073,9 +16074,16 @@ class Router: tier_param_keys: Final = ( tuple(pre_routing_hook_response.litellm_params or ()) if pre_routing_hook_response is not None else () ) - newly_forwarded: Final = tuple( - (key, value) for key, value in marker_params if key not in request_kwargs and key not in tier_param_keys + overridden_marker_keys: Final = frozenset(tier_param_keys) | ( + OUTPUT_TOKEN_CEILING_PARAMS if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_param_keys) else frozenset() ) + newly_forwarded: Final = tuple( + (key, value) + for key, value in marker_params + if key not in request_kwargs and key not in overridden_marker_keys + ) + if pre_routing_hook_response is not None: + self._stamp_client_output_ceiling(request_kwargs) request_kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, None) request_kwargs.update(newly_forwarded) if newly_forwarded: diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 7567a0ef37e..08d39841b61 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -107,6 +107,20 @@ Session affinity, user-turn classification, adaptive routing, routing plugins, c When cache costs change the model, the routing decision reports `cause: prompt_cache_cost`. Its signals include the original model, classification cause and both estimated costs +### Output token limits + +Set `max_tokens_from_tier_model` under `complexity_router_config` to choose how a routed request's output limit is determined: + +| Setting | Effective output limit | +| --- | --- | +| `true` (default) | Replace the caller's limit with the selected model group's ceiling | +| `"clamp"` | Use the smaller of the caller's positive limit and the selected model group's ceiling | +| `false` | Keep the caller's limit | + +Clamp reads `max_tokens`, `max_completion_tokens` and `max_output_tokens`. If several are supplied, it uses the smallest positive value. With no positive caller limit, it uses the model ceiling. For example, a caller requesting 100 tokens keeps that limit on a model with a 32,768-token ceiling; a request for 100,000 is capped at 32,768. Rerouting uses the original caller budget, and cache-aware cost estimates use the same effective limit. + +All modes preserve explicit output limits in a tier's own `litellm_params`. An unknown ceiling leaves the caller's value unchanged. In the dashboard, select **Respect caller budget** under **Advanced settings > Compatibility > Output token limit** to enable clamp. + ### Capability forecasting Set `classifier_type: capability` to use diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 12968bf7d4a..f5fd4ae9c98 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -23,6 +23,7 @@ import random import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from functools import partial from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType @@ -161,6 +162,17 @@ def _tier_name(tier: ComplexityTier | str) -> str: return tier.value if isinstance(tier, ComplexityTier) else tier +def _callers_output_cap(request_kwargs: Mapping[str, object]) -> int | None: + """The smallest positive output ceiling the caller sent under any of its three names, read + from the stamp when an earlier routing pass already rewrote the kwargs.""" + from litellm.router import as_output_cap, stamped_client_ceiling + + stamped: Final = stamped_client_ceiling(request_kwargs) + carriers: Final = request_kwargs if stamped is None else stamped + caps: Final = tuple(as_output_cap(carriers.get(carrier)) for carrier in OUTPUT_TOKEN_CEILING_PARAMS) + return min((cap for cap in caps if cap is not None and cap > 0), default=None) + + def _built_in_tier_or_none(tier_name: str) -> ComplexityTier | None: """The built-in tier a `tiers` key names, or None when the key is an operator-defined name.""" return ComplexityTier.__members__.get(tier_name) @@ -3028,7 +3040,9 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"No model configured for tier {tier_key} and no default_model set") - def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]: + def _litellm_params_for_model( + self, tier: ComplexityTier | str | None, model: str, request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else () entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None) explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({}) @@ -3037,7 +3051,12 @@ class ComplexityRouter(CustomLogger): ceiling: Final = self._group_output_ceiling(model) if ceiling is None: return explicit - return MappingProxyType({**explicit, "max_tokens": ceiling}) + callers_cap: Final = ( + _callers_output_cap(request_kwargs) if self.config.max_tokens_from_tier_model == "clamp" else None + ) + return MappingProxyType( + {**explicit, "max_tokens": ceiling if callers_cap is None else min(callers_cap, ceiling)} + ) @staticmethod def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str: @@ -3768,7 +3787,7 @@ class ComplexityRouter(CustomLogger): response: PreRoutingHookResponse, messages: list[dict[str, object]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick resolved_messages: Sequence[Mapping[str, object]] | None, - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: @@ -3865,7 +3884,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing=bool(decision.get("conversation_continuing", True)) if decision is not None else True, - tier_litellm_params=self._litellm_params_for_model(new_tier, new_model), + tier_litellm_params=self._litellm_params_for_model(new_tier, new_model, request_kwargs), context_escalation_original_tier=( decision.get("context_escalation_original_tier") if decision is not None else None ), @@ -3876,7 +3895,7 @@ class ComplexityRouter(CustomLogger): return HookResponse( model=new_model, messages=response.messages, - litellm_params=self._litellm_params_for_model(new_tier, new_model), + litellm_params=self._litellm_params_for_model(new_tier, new_model, request_kwargs), routing_decision=new_decision, ) @@ -3904,7 +3923,7 @@ class ComplexityRouter(CustomLogger): model_name: str, messages: list[dict[str, Any]] | None, # mutable-ok: forwarded verbatim to the router's own probe input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives ) -> bool: """Whether the router would find a deployment for this group ON THIS REQUEST. @@ -3962,7 +3981,7 @@ class ComplexityRouter(CustomLogger): messages: list[dict[str, object]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim resolved_messages: Sequence[Mapping[str, object]] | None, - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: @@ -4045,14 +4064,14 @@ class ComplexityRouter(CustomLogger): classifier_model=decision.get("classifier_model"), classifier_cost=decision.get("classifier_cost"), conversation_continuing=bool(decision.get("conversation_continuing", True)), - tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model), + tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model, request_kwargs), context_escalation_original_tier=decision.get("context_escalation_original_tier"), previous_decision=decision, ) return response.model_copy( update={ "model": new_model, - "litellm_params": self._litellm_params_for_model(candidate_tier, new_model), + "litellm_params": self._litellm_params_for_model(candidate_tier, new_model, request_kwargs), "routing_decision": new_decision, } ) @@ -4090,14 +4109,14 @@ class ComplexityRouter(CustomLogger): classifier_model=decision.get("classifier_model"), classifier_cost=decision.get("classifier_cost"), conversation_continuing=bool(decision.get("conversation_continuing", True)), - tier_litellm_params=self._litellm_params_for_model(None, default_model), + tier_litellm_params=self._litellm_params_for_model(None, default_model, request_kwargs), context_escalation_original_tier=decision.get("context_escalation_original_tier"), previous_decision=decision, ) return response.model_copy( update={ "model": default_model, - "litellm_params": self._litellm_params_for_model(None, default_model), + "litellm_params": self._litellm_params_for_model(None, default_model, request_kwargs), "routing_decision": default_decision, } ) @@ -4580,7 +4599,9 @@ class ComplexityRouter(CustomLogger): if pin_placement is not None and pin_context_original_tier is not None else pin_source_tier ) - session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) + session_tier_litellm_params: Final = self._litellm_params_for_model( + routed_pin_tier, routed_model, request_kwargs + ) has_original_messages: Final = messages is not None and len(messages) > 0 return self._with_session_deployment_affinity( await self._gate_response_health( @@ -4678,14 +4699,14 @@ class ComplexityRouter(CustomLogger): choice: Final = await choose_cached_model( router=self.litellm_router_instance, config=self.config, - params_for_model=self._litellm_params_for_model, + params_for_model=partial(self._litellm_params_for_model, request_kwargs=request_kwargs), response=response, request_kwargs=request_kwargs, messages=messages, ) if choice is None or not context_fit.accepts(choice.model): return response - params: Final = self._litellm_params_for_model(choice.tier, choice.model) + params: Final = self._litellm_params_for_model(choice.tier, choice.model, request_kwargs) decision: Final[StandardLoggingRoutingDecision] = { **response.routing_decision, "routed_model": choice.model, @@ -4714,7 +4735,7 @@ class ComplexityRouter(CustomLogger): async def _classify_and_route( self, model: str, - request_kwargs: dict, + request_kwargs: dict[str, object], messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, @@ -4772,7 +4793,7 @@ class ComplexityRouter(CustomLogger): ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM - default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model) + default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model, request_kwargs) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -4810,7 +4831,7 @@ class ComplexityRouter(CustomLogger): _tier_name(plan_floor), routed_model, ) - plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model) + plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model, request_kwargs) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -4839,7 +4860,9 @@ class ComplexityRouter(CustomLogger): ) keyword_plan_floored: Final = routed_tier != escalated_tier routed_model = await self._pick_model_for_tier(routed_tier, messages, resolved_messages, request_kwargs) - keyword_tier_litellm_params: Final = self._litellm_params_for_model(routed_tier, routed_model) + keyword_tier_litellm_params: Final = self._litellm_params_for_model( + routed_tier, routed_model, request_kwargs + ) keyword_cause: Final[RoutingDecisionCause] = ( "plan_mode" if keyword_plan_floored @@ -4919,7 +4942,7 @@ class ComplexityRouter(CustomLogger): outcome.signals, fallback_model, ) - fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model) + fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model, request_kwargs) return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, @@ -5001,7 +5024,7 @@ class ComplexityRouter(CustomLogger): routed_model, ) - tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model) + tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model, request_kwargs) classifier_model: Final = ( f"{outcome.jev_verdict.provider}/{outcome.jev_verdict.model}" if outcome.cause == "jev_classifier" and outcome.jev_verdict is not None diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index ec960a19f79..020f2d33742 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1462,7 +1462,7 @@ class ComplexityRouterConfig(LiteLLMBaseModel): "wording the built-ins don't cover, or after a client release changes its strings." ), ) - max_tokens_from_tier_model: bool = Field( + max_tokens_from_tier_model: bool | Literal["clamp"] = Field( default=True, description=( "Set max_tokens on every routed request to the output ceiling of the tier model it " @@ -1473,7 +1473,8 @@ class ComplexityRouterConfig(LiteLLMBaseModel): "from each deployment's model_info and then the model cost map; a tier model with a " "deployment whose ceiling is unknown keeps the caller's value. A max_tokens, " "max_completion_tokens or max_output_tokens in the tier's own litellm_params still " - "wins. Set false to forward the caller's value unchanged." + 'wins. Set "clamp" to keep a smaller positive caller value and cap a larger one at ' + "the ceiling. Set false to forward the caller's value unchanged." ), ) route_housekeeping_to_cheapest_tier: bool = Field( diff --git a/tests/unit/proxy/common_utils/test_cache_aware_routing.py b/tests/unit/proxy/common_utils/test_cache_aware_routing.py index 24fc1f22a6a..4eb22b2e34f 100644 --- a/tests/unit/proxy/common_utils/test_cache_aware_routing.py +++ b/tests/unit/proxy/common_utils/test_cache_aware_routing.py @@ -1,5 +1,6 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial from typing import Final import pytest @@ -148,7 +149,7 @@ async def _select( return await select_cached_model( router=router, config=config, - params_for_model=complexity._litellm_params_for_model, + params_for_model=partial(complexity._litellm_params_for_model, request_kwargs=request_kwargs), response=response, body=body, request_kwargs=request_kwargs, @@ -234,7 +235,9 @@ async def test_disabled_setting_does_not_access_prediction_services() -> None: await choose_cached_model( router=_router(), config=config, - params_for_model=ComplexityRouter("smart", _router(), config.model_dump())._litellm_params_for_model, + params_for_model=partial( + ComplexityRouter("smart", _router(), config.model_dump())._litellm_params_for_model, request_kwargs={} + ), response=_response(), request_kwargs={}, messages=None, @@ -348,6 +351,7 @@ def test_repeated_model_in_multiple_tiers_is_only_considered_once() -> None: (True, "custom_endpoint", "cheap"), (True, "compaction", "cheap"), (True, "guardrail", "cheap"), + (True, "clamp", "strong"), ], ) async def test_router_applies_opt_in_and_preserves_failure_semantics( @@ -368,7 +372,9 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( from litellm.router_strategy.complexity_router.context_compaction import initialize_compaction_state config: Final = _config( - cache_aware_routing=enabled, cache_aware_routing_timeout_ms=1 if behavior == "deadline" else 2000 + cache_aware_routing=enabled, + cache_aware_routing_timeout_ms=1 if behavior == "deadline" else 2000, + max_tokens_from_tier_model="clamp" if behavior == "clamp" else True, ) models: Final = _router( strong_output_rate=0.000048 if behavior == "tier_cost" else 0.000015, @@ -423,6 +429,7 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( body: Final = { **_body(), **({"max_tokens": 1000} if behavior in ("tier_cost", "tier_context") else {}), + **({"max_tokens": 100} if behavior == "clamp" else {}), **({"thinking": {"type": "enabled", "budget_tokens": 10000}} if behavior == "unsupported_shape" else {}), } kwargs: Final = { @@ -437,6 +444,7 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( else {} ), **({"guardrails": ["test-guardrail"]} if behavior == "guardrail" else {}), + **({"max_tokens": 100} if behavior == "clamp" else {}), } if expected is None: with pytest.raises(asyncio.CancelledError): @@ -458,6 +466,10 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( if enabled and behavior == "success": assert requests.qsize() == 4 assert response.routing_decision is not None + if behavior == "clamp": + assert "cache-aware:estimated-cost=0.01950000;original-cost=0.06390000" in ( + response.routing_decision["signals"] or () + ), "both arms must be estimated at the caller's 100 output tokens, not the 1024 default" assert response.routing_decision["cause"] == ( "prompt_cache_cost" if expected == "strong" else "heuristic_scorer" ) diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 2d27a025ffa..78dc2eaf314 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -36,11 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) -from litellm.router import as_output_cap +from litellm.router import as_output_cap, stamped_client_ceiling from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -13733,7 +13734,7 @@ def test_tier_model_params_are_used_by_pools_and_savings_baseline(mock_router_in assert router._tier_pools() == {"SIMPLE": ["mini"], "REASONING": ["opus", "abc"]} assert router._hardest_tier_models() == ("opus", "abc") - assert router._litellm_params_for_model(ComplexityTier.REASONING, "opus") == {"reasoning_effort": "xhigh"} + assert router._litellm_params_for_model(ComplexityTier.REASONING, "opus", {}) == {"reasoning_effort": "xhigh"} @pytest.mark.asyncio @@ -16917,9 +16918,10 @@ class TestMaxTokensFromTierModel: @staticmethod def _router( tier_litellm_params: dict | None = None, - max_tokens_from_tier_model: bool | None = None, + max_tokens_from_tier_model: bool | Literal["clamp"] | None = None, simple_deployments: list[dict] | None = None, extra_config: dict | None = None, + marker_params: Mapping[str, object] | None = None, ) -> Router: simple_tier: dict = {"model_name": "small"} if tier_litellm_params: @@ -16934,7 +16936,11 @@ class TestMaxTokensFromTierModel: model_list=[ { "model_name": "smart-router", - "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config}, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": config, + **(marker_params or {}), + }, }, *(simple_deployments or [TestMaxTokensFromTierModel.SMALL]), { @@ -17016,6 +17022,167 @@ class TestMaxTokensFromTierModel: assert sent["max_tokens"] == 8192 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "prompt, client_kwargs, expected", + [ + (COMPLEX_PROMPT, {"max_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + ("hi", {"max_tokens": 100000}, ("anthropic/claude-haiku-4-5", 8192)), + (COMPLEX_PROMPT, {}, ("anthropic/claude-sonnet-5", 64000)), + (COMPLEX_PROMPT, {"max_tokens": 0}, ("anthropic/claude-sonnet-5", 64000)), + (COMPLEX_PROMPT, {"max_completion_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + (COMPLEX_PROMPT, {"max_tokens": 3000, "max_completion_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + ], + ids=["below-ceiling", "above-ceiling", "omitted", "zero", "other-carrier", "smallest-carrier"], + ) + async def test_clamp_keeps_the_smaller_of_the_client_value_and_the_ceiling( + self, prompt: str, client_kwargs: dict[str, int], expected: tuple[str, int] + ) -> None: + sent: Final = await self._routed(self._router(max_tokens_from_tier_model="clamp"), prompt, **client_kwargs) + + assert (sent["model"], sent["max_tokens"]) == expected + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {"max_tokens"}) & sent.keys() + + @pytest.mark.asyncio + async def test_clamp_on_the_responses_surface_keeps_its_own_name(self) -> None: + router: Final = self._router(max_tokens_from_tier_model="clamp") + + sent: Final = await self._routed_responses(router, self.COMPLEX_PROMPT, max_output_tokens=2000) + + assert (sent["model"], sent["max_output_tokens"]) == ("anthropic/claude-sonnet-5", 2000) + assert "max_tokens" not in sent + + @pytest.mark.asyncio + @pytest.mark.parametrize("tier_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("surface", ["chat", "responses"]) + @pytest.mark.parametrize("caller_cap", [None, 100]) + async def test_clamp_leaves_the_operators_own_tier_ceiling_in_charge( + self, + tier_carrier: str, + marker_carrier: str, + surface: str, + caller_cap: int | None, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + surface_key: Final = "max_output_tokens" if surface == "responses" else "max_tokens" + caller_key: Final = "max_tokens" if surface == "responses" else "max_output_tokens" + request_caps: Final = {caller_key: caller_cap} if caller_cap is not None else {} + router: Final = self._router( + tier_litellm_params={tier_carrier: 4321}, + max_tokens_from_tier_model="clamp", + marker_params={marker_carrier: 2000}, + simple_deployments=[{ + "model_name": "small", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-only", + "api_base": "https://tier.test/v1", + surface_key: 9000, + }, + "model_info": {"max_output_tokens": 16384}, + }], + ) + + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host="tier.test").mock(side_effect=TestHealthFallbackDispatch._http_response) + if surface == "responses": + await router.aresponses(model="smart-router", input="hi", **request_caps) + else: + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}], **request_caps + ) + + assert len(upstream.calls) == 1 + sent: Final = json.loads(upstream.calls[0].request.content) + wire_key: Final = "max_output_tokens" if surface == "responses" else "max_completion_tokens" + assert {key: value for key, value in sent.items() if key in OUTPUT_TOKEN_CEILING_PARAMS} == {wire_key: 4321} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "callers_original, expected", + [({"max_tokens": 2000}, 2000), ({}, 64000)], + ids=["caller-sent-a-value", "caller-sent-nothing"], + ) + async def test_clamp_on_a_rerouted_request_uses_the_callers_original_value( + self, callers_original: dict[str, int], expected: int + ) -> None: + """A retry re-enters routing with kwargs an earlier pass already rewrote; the stamp holds + what the caller actually sent.""" + router: Final = self._router(max_tokens_from_tier_model="clamp") + rewritten: Final[dict[str, object]] = { + "max_tokens": 8192, + "metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: callers_original}, + } + + sent: Final = await self._routed(router, self.COMPLEX_PROMPT, **rewritten) + + assert sent["max_tokens"] == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("caller_cap", [None, 3000]) + @pytest.mark.parametrize("tier_cap", [None, 4321]) + @pytest.mark.parametrize("known_ceiling", [True, False]) + async def test_rerouting_preserves_caller_budget_without_adopting_marker_defaults( + self, marker_carrier: str, caller_cap: int | None, tier_cap: int | None, known_ceiling: bool + ) -> None: + router: Final = self._router( + max_tokens_from_tier_model="clamp", + marker_params={marker_carrier: 2000}, + tier_litellm_params={"max_tokens": tier_cap} if tier_cap else None, + simple_deployments=None if known_ceiling else [{ + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + }], + ) + request: Final[dict[str, object]] = {} if caller_cap is None else {"max_output_tokens": caller_cap} + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] + ) + first: Final = {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, + messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], + ) + + expected_first: Final = ( + {"max_tokens": tier_cap or caller_cap or 8192} + if known_ceiling or tier_cap else + {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})} + ) + assert first == expected_first + assert request["max_tokens"] == (caller_cap or 64000) + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {"max_tokens"}) & request.keys() + + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] + ) + assert {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} == first + + @pytest.mark.asyncio + @pytest.mark.parametrize("mode", [False, "clamp"]) + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + async def test_unpinned_output_limits_still_use_marker_defaults( + self, mode: bool | Literal["clamp"], marker_carrier: str + ) -> None: + router: Final = self._router( + max_tokens_from_tier_model=mode, + marker_params={marker_carrier: 2000}, + simple_deployments=[{ + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + }], + ) + sent: Final = await self._routed(router) + + assert {key: sent[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in sent} == {marker_carrier: 2000} + + def test_an_unknown_mode_is_rejected(self) -> None: + with pytest.raises(ValidationError): + ComplexityRouterConfig(tiers={"SIMPLE": "small"}, max_tokens_from_tier_model="sometimes") + @pytest.mark.asyncio async def test_a_tier_model_with_an_unknown_ceiling_keeps_the_client_value(self): unmapped: dict = {"model_name": "small", "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}} @@ -17112,7 +17279,7 @@ class TestMaxTokensFromTierModel: complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"}, ) - assert dict(strategy._litellm_params_for_model(None, "big")) == {"max_tokens": 64000} + assert dict(strategy._litellm_params_for_model(None, "big", {})) == {"max_tokens": 64000} @pytest.mark.asyncio async def test_a_fallback_into_a_plain_group_gets_the_callers_ceiling_back(self): @@ -17228,29 +17395,52 @@ class TestMaxTokensFromTierModel: def test_a_client_cap_is_read_as_an_integer_or_ignored(self, value, expected): assert as_output_cap(value) == expected + @pytest.mark.parametrize( + "request_kwargs, expected", + [ + ({"max_tokens": 8192}, None), + ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {}}}, {}), + ({"litellm_metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_tokens": "2000"}}}, {"max_tokens": 2000}), + ( + {"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True}}}, + {"max_output_tokens": 5}, + ), + ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: "not-a-dict"}}, None), + ], + ids=["no-stamp", "empty-stamp", "litellm-metadata-bucket", "only-integer-carriers", "malformed-stamp"], + ) + def test_the_stamp_reads_back_only_the_callers_integer_carriers( + self, request_kwargs: dict[str, object], expected: dict[str, int] | None + ) -> None: + stamped: Final = stamped_client_ceiling(request_kwargs) + + assert (None if stamped is None else dict(stamped)) == expected + def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self): stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}} - Router._restore_client_ceiling_no_tier_pins(stamped) + Router._restore_client_output_ceiling(stamped) assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192} coerced: dict = { "max_tokens": 64000, "metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}}, } - Router._restore_client_ceiling_no_tier_pins(coerced) + Router._restore_client_output_ceiling(coerced) assert {k: v for k, v in coerced.items() if k != "metadata"} == { "max_tokens": 8192, "max_completion_tokens": 100, } unstamped: dict = {"max_tokens": 64000, "metadata": {}} - Router._restore_client_ceiling_no_tier_pins(unstamped) + Router._restore_client_output_ceiling(unstamped) assert unstamped["max_tokens"] == 64000 @pytest.mark.asyncio async def test_pinning_stamps_the_callers_carriers_once(self): router = self._router() request_kwargs: dict = {"max_completion_tokens": 8192} + router._stamp_client_output_ceiling(request_kwargs) + assert stamped_client_ceiling(request_kwargs) == {"max_completion_tokens": 8192} first = router._pin_tier_params_onto_request( model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 72a08c53267..b59aa051e38 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -452,7 +452,7 @@ export interface ComplexityRouterConfigValue { route_housekeeping_to_cheapest_tier?: boolean; housekeeping_patterns?: string[]; reminder_markers?: ReminderMarkerPair[]; - max_tokens_from_tier_model?: boolean; + max_tokens_from_tier_model?: boolean | "clamp"; classifier_plugin_timeout_ms?: number; } diff --git a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx index 6e7b1489a4c..f596ecfdc62 100644 --- a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx @@ -1,7 +1,28 @@ import { Switch } from "@/components/ui/switch"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import React from "react"; import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +const OUTPUT_TOKEN_POLICIES = [ + { + value: true, + label: "Use model ceiling", + description: "Replace the caller's output limit with the selected model's ceiling. This is the default.", + }, + { + value: "clamp", + label: "Respect caller budget", + description: + "Use the smaller of the caller's output limit and the model ceiling. Use the ceiling when no limit is sent.", + }, + { + value: false, + label: "Pass through caller limit", + description: "Forward the caller's output limit unchanged.", + }, +] as const; + const ResponseFormatControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; @@ -18,18 +39,31 @@ const ResponseFormatControls: React.FC<{ Return the resolved underlying model name in responses instead of the autorouter alias. -
- onChange({ ...value, max_tokens_from_tier_model: enabled })} - aria-label="Cap max_tokens at the tier model's output ceiling" - /> - Cap max_tokens at the tier model's output ceiling +
+ Output token limit + { + if (typeof mode === "boolean" || mode === "clamp") { + onChange({ ...value, max_tokens_from_tier_model: mode }); + } + }} + aria-label="Output token limit" + > + {OUTPUT_TOKEN_POLICIES.map((policy) => ( + + ))} + + + Explicit tier output limits take precedence. Unknown model ceilings keep the caller's limit. +
- - Replace the caller's max_tokens with the routed tier model's output ceiling so one client value fits - every tier. Off forwards the caller's value unchanged. - ); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 89d419d0487..b25baa0d1e4 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -1315,6 +1315,17 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { expect(payload.fallback_tier).toBe("CASUAL"); }); + it("keeps clamp and explicit output limits on custom tiers", () => { + const payload = build({ + maxTokensFromTierModel: "clamp", + tierModelParams: { sec: { "o1-preview": { max_completion_tokens: 700 } } }, + }); + expect(payload.max_tokens_from_tier_model).toBe("clamp"); + expect(payload.tier_model_configs?.SECURITY_REVIEW).toEqual([ + { model_name: "o1-preview", litellm_params: { max_completion_tokens: 700 } }, + ]); + }); + it("forces the LLM classifier even when the form still holds heuristic, which the backend rejects", () => { expect(build({ classifierType: "heuristic" }).classifier_type).toBe("llm"); }); @@ -1595,6 +1606,15 @@ describe("classifier vision wire payload", () => { }); describe("advanced complexity router fields", () => { + it.each([undefined, true, false, "clamp"] as const)("serializes output token policy %s", (mode) => { + const payload = buildComplexityRouterConfig({ ...baseParams, maxTokensFromTierModel: mode }); + if (mode === false || mode === "clamp") { + expect(payload.max_tokens_from_tier_model).toBe(mode); + } else { + expect(payload).not.toHaveProperty("max_tokens_from_tier_model"); + } + }); + it("normalizes lists, reminder markers, and explicit false values", () => { const payload = buildComplexityRouterConfig({ ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 09bf0ea41a9..511a1653aae 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -226,7 +226,7 @@ export interface BuildComplexityRouterConfigParams { routeHousekeepingToCheapestTier?: boolean; housekeepingPatterns?: string[]; reminderMarkers?: ReminderMarkerPair[]; - maxTokensFromTierModel?: boolean; + maxTokensFromTierModel?: ComplexityRouterConfigValue["max_tokens_from_tier_model"]; classifierPluginTimeoutMs?: number; } @@ -310,7 +310,7 @@ export interface ComplexityRouterConfigPayload { route_housekeeping_to_cheapest_tier?: boolean; housekeeping_patterns?: string[]; reminder_markers?: ReminderMarkerPair[]; - max_tokens_from_tier_model?: boolean; + max_tokens_from_tier_model?: ComplexityRouterConfigValue["max_tokens_from_tier_model"]; classifier_plugin_timeout_ms?: number; } @@ -827,7 +827,9 @@ export const buildComplexityRouterConfig = ({ ...cleanedLists, ...(routeHousekeepingToCheapestTier === false && { route_housekeeping_to_cheapest_tier: false }), ...(cleanedReminderMarkers && cleanedReminderMarkers.length > 0 && { reminder_markers: cleanedReminderMarkers }), - ...(maxTokensFromTierModel === false && { max_tokens_from_tier_model: false }), + ...((maxTokensFromTierModel === false || maxTokensFromTierModel === "clamp") && { + max_tokens_from_tier_model: maxTokensFromTierModel, + }), ...(hasValidCustomClassifierTimeout && { classifier_plugin_timeout_ms: classifierPluginTimeoutMs }), ...scorerKnobs, }; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 9445e18913e..6238e30ce07 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -47,6 +47,29 @@ const hydratedState: KeywordMatchingState = { matchThreshold: 0.72, }; +describe("output token policy edits", () => { + it.each([undefined, true, false, "clamp"] as const)("preserves and edits stored mode %s", (mode) => { + const stored = { + ...STORED, + ...(mode !== undefined && { max_tokens_from_tier_model: mode }), + tier_model_configs: { SIMPLE: [{ model_name: "gpt-4o-mini", litellm_params: { max_tokens: 700 } }] }, + }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + expect(hydrated.max_tokens_from_tier_model).toBe(mode); + const unchanged = buildUpdatedComplexityRouterConfig(stored, { ...hydrated, return_raw_model_name: true }); + expect(unchanged.max_tokens_from_tier_model).toBe(mode); + expect("max_tokens_from_tier_model" in unchanged).toBe(mode !== undefined); + for (const nextMode of [true, false, "clamp"] as const) { + const saved = buildUpdatedComplexityRouterConfig(stored, { + ...hydrated, + max_tokens_from_tier_model: nextMode, + }); + expect(saved.max_tokens_from_tier_model).toBe(nextMode === true && mode !== true ? undefined : nextMode); + expect(saved.tier_model_configs).toEqual(stored.tier_model_configs); + } + }); +}); + describe("buildUpdatedComplexityRouterConfig keyword matching", () => { it.each([undefined, false, true])("preserves cache settings through edit and save: %s", (enabled) => { const stored = { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx index 72d557d1f8f..e8c9288cfc5 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx @@ -472,9 +472,17 @@ describe("EditAutoRouterModal advanced field round trips", () => { modelPatchUpdateCall.mockClear(); }); - it("hydrates housekeeping and reminder fields, then omits the default max-token value after editing", async () => { + it("reopens clamp and saves the selected output policy without changing other advanced settings", async () => { const user = userEvent.setup(); - renderAdvancedModal(); + renderAdvancedModal({ + modelData: { + ...MODEL_DATA, + litellm_params: { + ...MODEL_DATA.litellm_params, + complexity_router_config: { ...storedAdvancedConfig, max_tokens_from_tier_model: "clamp" }, + }, + }, + }); openAutoRouterAdvanced("Housekeeping Routing"); expect(screen.getByRole("switch", { name: "Route housekeeping calls to the cheapest tier" })).not.toBeChecked(); @@ -485,13 +493,17 @@ describe("EditAutoRouterModal advanced field round trips", () => { expect(screen.getByLabelText("Closing tag")).toHaveValue(""); openAutoRouterAdvanced("Response Format"); - const maxTokensSwitch = screen.getByRole("switch", { name: "Cap max_tokens at the tier model's output ceiling" }); - await user.click(maxTokensSwitch); + expect(screen.getByRole("radio", { name: /^Respect caller budget/ })).toBeChecked(); + await user.click(screen.getByRole("radio", { name: /^Pass through caller limit/ })); + expect(screen.getByRole("radio", { name: /^Pass through caller limit/ })).toBeChecked(); + await user.click(screen.getByRole("radio", { name: /^Use model ceiling/ })); + expect(screen.getByRole("radio", { name: /^Use model ceiling/ })).toBeChecked(); + await user.click(screen.getByRole("radio", { name: /^Respect caller budget/ })); await waitFor(() => expect(screen.getByRole("button", { name: /save changes/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); - expect(savedConfig()).not.toHaveProperty("max_tokens_from_tier_model"); + expect(savedConfig().max_tokens_from_tier_model).toBe("clamp"); expect(savedConfig()).toMatchObject({ route_housekeeping_to_cheapest_tier: false, housekeeping_patterns: ["conversation title"], diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts index 233b5aed922..75fd6e15bc1 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts @@ -184,7 +184,8 @@ export const hydrateComplexityRouterConfig = ( ? parsedConfig.reminder_markers.filter(isReminderMarkerPair) : undefined, max_tokens_from_tier_model: - typeof parsedConfig.max_tokens_from_tier_model === "boolean" + typeof parsedConfig.max_tokens_from_tier_model === "boolean" || + parsedConfig.max_tokens_from_tier_model === "clamp" ? parsedConfig.max_tokens_from_tier_model : undefined, classifier_plugin_timeout_ms: diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index b398780d14e..f4e45e473df 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -877,6 +877,13 @@ describe("autorouter_presets", () => { }); describe("buildPresetPrefill", () => { + it.each([undefined, true, false, "clamp"] as const)("preserves a preset's output token policy %s", (mode) => { + const preset = getPresetByKey("anthropic_family")!; + const config = { ...preset.complexity_router_config, max_tokens_from_tier_model: mode }; + const prefill = buildPresetPrefill(config, groupsOnly(getRequiredModelsInPreset(preset))); + expect(prefill.complexityRouterConfig.max_tokens_from_tier_model).toBe(mode); + }); + it("preserves Laya settings and drops inactive classifier settings when prefilling", () => { const config = { tiers: { SIMPLE: ["fast"], MEDIUM: [], COMPLEX: [], REASONING: [] }, diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 0a83a06050b..fd400a8163a 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -364,6 +364,7 @@ export const buildPresetPrefill = ( tier_distance_penalty: config.tier_distance_penalty, adaptive_eligible: config.adaptive_eligible, return_raw_model_name: config.return_raw_model_name, + max_tokens_from_tier_model: config.max_tokens_from_tier_model, dimension_weights: hydrateDimensionWeights(config.dimension_weights), custom_dimensions: hydrateCustomDimensions(config.custom_dimensions), tier_boundaries: hydrateTierBoundaries(config.tier_boundaries), diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d03f4b4cffd..e3f80b37498 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -42620,10 +42620,10 @@ export interface components { match_threshold: number; /** * Max Tokens From Tier Model - * @description Set max_tokens on every routed request to the output ceiling of the tier model it lands on, replacing whatever the caller sent. A caller behind an auto-router cannot pick one value that fits every tier: the smallest tier's ceiling starves a bigger tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The ceiling is the smallest max_output_tokens across the tier model's deployments, read from each deployment's model_info and then the model cost map; a tier model with a deployment whose ceiling is unknown keeps the caller's value. A max_tokens, max_completion_tokens or max_output_tokens in the tier's own litellm_params still wins. Set false to forward the caller's value unchanged. + * @description Set max_tokens on every routed request to the output ceiling of the tier model it lands on, replacing whatever the caller sent. A caller behind an auto-router cannot pick one value that fits every tier: the smallest tier's ceiling starves a bigger tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The ceiling is the smallest max_output_tokens across the tier model's deployments, read from each deployment's model_info and then the model cost map; a tier model with a deployment whose ceiling is unknown keeps the caller's value. A max_tokens, max_completion_tokens or max_output_tokens in the tier's own litellm_params still wins. Set "clamp" to keep a smaller positive caller value and cap a larger one at the ceiling. Set false to forward the caller's value unchanged. * @default true */ - max_tokens_from_tier_model: boolean; + max_tokens_from_tier_model: boolean | "clamp"; /** * Modality Pin Override * @description Let modality_routing replace a kept session-affinity pin on the turns that carry an image. Without this, a session pinned to a text-only model fails every image turn with a provider 400, since the pin is exempt from the modality gate. When enabled, such a turn routes to a capable model for that request only and the stored pin is left untouched, so the next text turn replays the session's own model; the override is reported as cause modality_pin_override and is never itself pinned. Inert unless modality_routing is also enabled. From c23cff7e459aa83373209e273179481227155005 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:30:48 -0400 Subject: [PATCH 021/147] fix(mcp): enforce deny-first tool policy consistently (#48625) --- .../proxy/mcp_server/mcp_server_manager.py | 25 ++++- litellm/proxy/mcp_server/operations.py | 27 ++--- litellm/proxy/mcp_server/utils.py | 9 ++ .../mcp_server/test_mcp_server_manager.py | 102 ++++++++++++------ .../test_mcp_server_tool_calls_and_headers.py | 35 +++--- 5 files changed, 129 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 8e613a8a6e8..f7146ac9428 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -1705,6 +1705,24 @@ def warn_on_shared_identifier_prefixes(servers: Iterable[MCPServer]) -> None: ) +def _warn_overlapping_tool_policies(server: MCPServer) -> None: + overlapping_tools: Final = tuple( + entry + for entry in server.allowed_tools or () + if any( + match_known_tool_name(name, server, server.disallowed_tools or ()) is not None + for name in (entry, strip_known_server_prefix(entry, server)) + ) + ) + if overlapping_tools: + verbose_logger.warning( + "MCP server %s has entries in both allowed_tools and disallowed_tools: %s. " + "disallowed_tools takes precedence; matching tools will be blocked.", + server.name, + ", ".join(overlapping_tools), + ) + + def _warn_legacy_delegate_auth_if_applicable(server: MCPServer, *, source: str) -> None: """Direct legacy delegated OAuth configurations to the admitted replacement.""" if server.auth_type != MCPAuth.oauth2: @@ -2854,6 +2872,7 @@ class MCPServerManager: oauth_identity_binding=server_config.get("oauth_identity_binding", None), ) self.assign_unique_short_prefix(new_server) + _warn_overlapping_tool_policies(new_server) _warn_legacy_delegate_auth_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) self.invalidate_server_definition_caches(server_id) @@ -5943,12 +5962,10 @@ class MCPServerManager: cut a second boundary out of a native name that opens with the server prefix. """ from litellm.proxy.mcp_server.utils import ( - server_applies_tool_allowlist, + is_tool_allowed_for_server, ) - if server_applies_tool_allowlist(server): - return match_known_tool_name(tool_name, server, server.allowed_tools or ()) is not None - return match_known_tool_name(tool_name, server, server.disallowed_tools or ()) is None + return is_tool_allowed_for_server(tool_name, server) def validate_allowed_params(self, tool_name: str, arguments: _ToolArguments, server: MCPServer) -> None: """ diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index 763e3578a96..4c2cbe2dfbf 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -800,28 +800,17 @@ def filter_tools_by_allowed_tools( Filtered list of tools """ from litellm.proxy.mcp_server.utils import ( + is_tool_allowed_for_server, server_applies_tool_allowlist, ) - tools_to_return = tools - - # Filter by allowed_tools (whitelist) - if server_applies_tool_allowlist(mcp_server): - if not mcp_server.allowed_tools: - return [] - tools_to_return = [ - tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools, mcp_server) - ] - - # Filter by disallowed_tools (blacklist) - if mcp_server.disallowed_tools: - tools_to_return = [ - tool - for tool in tools_to_return - if not _tool_name_matches(tool.name, mcp_server.disallowed_tools, mcp_server) - ] - - return tools_to_return + if not server_applies_tool_allowlist(mcp_server) and not mcp_server.disallowed_tools: + return tools + return [ + tool + for tool in tools + if is_tool_allowed_for_server(strip_known_server_prefix(tool.name, mcp_server), mcp_server) + ] def apply_display_name_overrides( diff --git a/litellm/proxy/mcp_server/utils.py b/litellm/proxy/mcp_server/utils.py index ee8f3ea3b30..3f72804ff35 100644 --- a/litellm/proxy/mcp_server/utils.py +++ b/litellm/proxy/mcp_server/utils.py @@ -266,6 +266,15 @@ def server_applies_tool_allowlist(mcp_server: object) -> bool: return is_server_tool_allowlist_enforced(mcp_server) or bool(allowed_tools) +def is_tool_allowed_for_server(tool_name: str, server: MCPServer) -> bool: + """Evaluate server tool policy for a bare name, with explicit denials taking precedence.""" + if match_known_tool_name(tool_name, server, server.disallowed_tools or ()) is not None: + return False + return not server_applies_tool_allowlist(server) or ( + match_known_tool_name(tool_name, server, server.allowed_tools or ()) is not None + ) + + def validate_and_normalize_mcp_server_payload(payload: McpServerPayloadLike) -> None: """ Validate and normalize MCP server payload fields (server_name, alias, and diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index c76326c51f6..c42579dc7f7 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -71,6 +71,7 @@ from litellm.proxy.mcp_server.legacy_callbacks import ElicitationCallback from litellm.proxy.mcp_server.mcp_context import active_mcp_request_ctx_var from litellm.proxy.mcp_server.mcp_server_manager import ( ListedToolsCaller, + MCPServerConfig, MCPServerManager, _deserialize_json_dict, _deserialize_json_list, @@ -1076,6 +1077,50 @@ class TestMCPServerManager: assert server.alias == "friendly_alias" assert server.server_name == "validserver" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "allowed_entry,denied_entry,expect_warning", + [ + ("read_status", "read_status", True), + ("policy-read_status", "read_status", True), + ("read_status", "policy-read_status", True), + ("policy-read_status", "policy-read_status", True), + ("read_status", "READ_STATUS", False), + ("read_status", "delete_records", False), + ], + ) + async def test_config_tool_policy_overlap_warns_without_changing_lists( + self, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + allowed_entry: str, + denied_entry: str, + expect_warning: bool, + ) -> None: + monkeypatch.delenv("LITELLM_ENABLE_MCP_STDIO", raising=False) + manager: Final = MCPServerManager() + config: Final[dict[str, MCPServerConfig]] = { + "policy_server": { + "alias": "policy", + "transport": MCPTransport.stdio, + "command": "python", + "allowed_tools": [allowed_entry], + "disallowed_tools": [denied_entry], + } + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + + server: Final = next(iter(manager.config_mcp_servers.values())) + warnings: Final = tuple( + message for message in caplog.messages if "both allowed_tools and disallowed_tools" in message + ) + assert len(warnings) == int(expect_warning) + assert all("will be blocked" in message for message in warnings) + assert server.allowed_tools == [allowed_entry] + assert server.disallowed_tools == [denied_entry] + def _oauth2_config(self, **overrides): base = { "url": "https://example.com/mcp", @@ -6431,53 +6476,42 @@ class TestMCPServerManager: ) @pytest.mark.asyncio - async def test_pre_call_tool_check_allowed_tools_takes_precedence(self): - """Test that allowed_tools list takes precedence over disallowed_tools list""" - manager = MCPServerManager() - - # Create server with both allowed_tools and disallowed_tools - # Note: The logic in check_allowed_or_banned_tools prioritizes allowed_tools - server = MCPServer( + @pytest.mark.parametrize("tool_name,denied", [("tool1", False), ("tool2", True), ("tool3", True)]) + async def test_pre_call_tool_check_disallowed_tools_takes_precedence( + self, tool_name: str, denied: bool + ) -> None: + manager: Final = MCPServerManager() + server: Final = MCPServer( server_id="test-server", name="test-server", transport=MCPTransport.stdio, allowed_tools=["tool1", "tool2"], - disallowed_tools=["tool2", "tool3"], # tool2 is in both lists + disallowed_tools=["tool2", "tool3"], ) - - # Mock dependencies - set object_permission and object_permission_id to None - # so permission checks return None (no restrictions) user_api_key_auth: Final = UserAPIKeyAuth() - proxy_logging_obj = _mock_proxy_logging() - - # Mock the async methods that pre_call_tool_check calls - proxy_logging_obj.create_mcp_request_object_from_kwargs = MagicMock(return_value={}) - proxy_logging_obj.convert_mcp_to_llm_format = MagicMock(return_value={}) - proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) - - # tool2 should be allowed since it's in allowed_tools (takes precedence) - await manager.pre_call_tool_check( - name="tool2", - arguments={"param": "value"}, - server_name="test-server", - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - server=server, - ) - - # tool3 should be blocked since it's not in allowed_tools - with pytest.raises(HTTPException) as exc_info: + if not denied: await manager.pre_call_tool_check( - name="tool3", - arguments={"param": "value"}, + name=tool_name, + arguments={}, server_name="test-server", user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=None, + server=server, + ) + return + + with pytest.raises(HTTPException) as exc_info: + await manager.pre_call_tool_check( + name=tool_name, + arguments={}, + server_name="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=None, server=server, ) assert exc_info.value.status_code == 403 - assert "Tool tool3 is not allowed for server test-server" in exc_info.value.detail["error"] + assert f"Tool {tool_name} is not allowed for server test-server" in exc_info.value.detail["error"] async def test_get_tools_from_server_add_prefix(self): """Verify _get_tools_from_server respects add_prefix True/False.""" diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index aa61deaf554..18907228b00 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -10941,7 +10941,7 @@ class TestListFiltersHonorThePrefixBoundary: assert _tool_name_matches("eiG-read_wiki_contents", ["deepwiki_cfg-read_wiki_contents"], server) - def test_discovery_hides_exactly_what_dispatch_refuses(self, monkeypatch): + def test_discovery_hides_exactly_what_dispatch_refuses(self, monkeypatch: pytest.MonkeyPatch) -> None: """Both halves of one decision, driven through both production paths. A spelling the blocklist enforces but the filter misses leaves a blocked @@ -10963,7 +10963,10 @@ class TestListFiltersHonorThePrefixBoundary: monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") - def _server(**overrides): + def _server( + allowed_tools: list[str] | None = None, + disallowed_tools: list[str] | None = None, + ) -> MCPServer: return MCPServer( server_id=self.SERVER_ID, name="deepwiki_prod", @@ -10972,27 +10975,35 @@ class TestListFiltersHonorThePrefixBoundary: short_prefix="eiG", url="http://127.0.0.1:5115/mcp", transport=MCPTransport.http, - **overrides, + allowed_tools=allowed_tools, + disallowed_tools=disallowed_tools, ) - manager = MCPServerManager() + manager: Final = MCPServerManager() manager._create_prefixed_tools( [MCPTool(name="read_wiki_contents", description="", inputSchema={"type": "object"})], _server(), ) - registered = sorted(manager.tool_name_to_mcp_server_name_mapping) + registered: Final = sorted(manager.tool_name_to_mcp_server_name_mapping) assert len(registered) > 1 - published = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"}) + published: Final = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"}) for spelling in registered: - for entry, expected in ((spelling, True), (spelling.upper(), False)): - server = _server(disallowed_tools=[entry]) + policies: Final[tuple[tuple[list[str] | None, list[str], bool], ...]] = ( + (None, [spelling], True), + (None, [spelling.upper()], False), + (["read_wiki_contents"], [spelling], True), + ([spelling], ["read_wiki_contents"], True), + ([spelling], [spelling.upper()], False), + ) + for allowed_tools, disallowed_tools, expected in policies: + server: Final = _server(allowed_tools, disallowed_tools) - refused = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) - hidden = filter_tools_by_allowed_tools([published], server) == [] + refused: Final = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) + hidden: Final = filter_tools_by_allowed_tools([published], server) == [] - assert refused == hidden, entry - assert refused is expected, entry + assert refused == hidden, (allowed_tools, disallowed_tools) + assert refused is expected, (allowed_tools, disallowed_tools) @pytest.mark.asyncio @pytest.mark.parametrize( From ad4b35d791bd79b142f25a022e9f8a77230f0798 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:46:15 -0400 Subject: [PATCH 022/147] fix(mcp): accept YAML OpenAPI documents (#48622) --- .../mcp_server/openapi_to_mcp_generator.py | 101 +++++++++++- .../mcp_server/test_mcp_server_manager.py | 37 ++++- .../mcp_server/test_openapi_spec_path_url.py | 152 +++++++++++++++++- 3 files changed, 284 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py index 96d178dbeb1..0fa904cb582 100644 --- a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py @@ -5,15 +5,29 @@ This module is used to generate MCP tools from OpenAPI specs. import asyncio import contextvars import json +import math import os import re from collections.abc import Mapping, Sequence from pathlib import PurePosixPath -from typing import Final, Literal, Protocol, TypeAlias, TypedDict, overload +from typing import ( + ClassVar, + Final, + Literal, + Protocol, + TypeAlias, + TypedDict, + cast, # noqa: TID251 # existing OpenAPI Protocol fields have no runtime schema + overload, +) from urllib.parse import quote import httpx +import yaml +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required +from yaml.events import AliasEvent, ScalarEvent +from yaml.nodes import MappingNode, Node, ScalarNode from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError from litellm.proxy.mcp_server.exceptions import ( @@ -30,6 +44,18 @@ _OPENAPI_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") OPENAPI_TOOL_NAME_MAX_LEN: Final = 128 _OPENAPI_TOOL_NAME_MAX_LEN: Final = OPENAPI_TOOL_NAME_MAX_LEN +_OPENAPI_DOCUMENT: Final = TypeAdapter(dict[str, object]) +_MAX_YAML_INT_LENGTH: Final = 1024 +# YAML 1.2 core scalars; isolate these from PyYAML's global YAML 1.1 resolver. +_YAML_SCALARS: Final = { + "tag:yaml.org,2002:null": re.compile(r"^(?:null|Null|NULL|~|)$"), + "tag:yaml.org,2002:bool": re.compile(r"^(?:true|True|TRUE|false|False|FALSE)$"), + "tag:yaml.org,2002:int": re.compile(r"^(?:[-+]?[0-9]+|0o[0-7]+|0x[0-9a-fA-F]+)$"), + "tag:yaml.org,2002:float": re.compile( + r"^(?:[-+]?(?:\.[0-9]+|[0-9]+(?:\.[0-9]*)?)(?:[eE][-+]?[0-9]+)?" + r"|[-+]?\.(?:inf|Inf|INF)|\.(?:nan|NaN|NAN))$" + ), +} def sanitize_openapi_tool_name(raw_name: str) -> str: @@ -199,6 +225,70 @@ def load_openapi_spec(filepath: str) -> _OpenAPISpec: return asyncio.run(load_openapi_spec_async(filepath)) +class _OpenAPISafeLoader(yaml.SafeLoader): + yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state + None: list(_YAML_SCALARS.items()) + } + + def compose_node(self, parent: Node | None, index: int) -> Node | None: + event: Final[object] = self.peek_event() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML event stubs + if isinstance(event, AliasEvent): + raise ValueError("YAML aliases are not supported in OpenAPI specs") # noqa: TRY004 # callers handle malformed documents as ValueError + node: Final = super().compose_node(parent, index) + # PyYAML otherwise implicitly resolves even an explicit nonspecific ! tag. + if isinstance(event, ScalarEvent) and event.tag == "!" and isinstance(node, ScalarNode): + node.tag = "tag:yaml.org,2002:str" + return node + + def construct_object(self, node: Node, deep: bool = False) -> object: + pattern: Final = _YAML_SCALARS.get(node.tag) + if pattern is not None: + if not isinstance(node, ScalarNode): + raise ValueError("Invalid YAML OpenAPI spec") + scalar: Final[str] = node.value + if node.tag == "tag:yaml.org,2002:int" and len(scalar) > _MAX_YAML_INT_LENGTH: + raise ValueError("YAML integer is too long in OpenAPI spec") + if pattern.fullmatch(scalar) is None: + raise ValueError("Invalid YAML OpenAPI spec") + if node.tag == "tag:yaml.org,2002:int": + # YAML 1.2 decimal values with leading zeroes are not octal. + return int(scalar, 0 if scalar.startswith(("0o", "0x")) else 10) + if node.tag == "tag:yaml.org,2002:float": + value: Final = self.construct_yaml_float(node) + if not math.isfinite(value): + raise ValueError("Invalid YAML OpenAPI spec") + return value + elif node.tag not in {"tag:yaml.org,2002:str", "tag:yaml.org,2002:seq", "tag:yaml.org,2002:map"} or ( + node.tag == "tag:yaml.org,2002:str" and not isinstance(node, ScalarNode) + ): + raise ValueError("Invalid YAML OpenAPI spec") + if isinstance(node, MappingNode): + pairs: Final[Sequence[tuple[Node, Node]]] = node.value + if any(not isinstance(key, ScalarNode) or key.tag != "tag:yaml.org,2002:str" for key, _ in pairs): + raise ValueError("Invalid YAML OpenAPI spec") + return super().construct_object( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML stubs + node, deep + ) + + +def _validate_openapi_spec(value: object) -> _OpenAPISpec: + try: + document: Final = _OPENAPI_DOCUMENT.validate_python(value) + except ValidationError as exc: + raise ValueError("OpenAPI spec must be a JSON or YAML object") from exc + return cast(_OpenAPISpec, document) # cast-ok: preserve existing lazy OpenAPI Protocol validation + + +def _parse_openapi_spec(content: str | bytes) -> _OpenAPISpec: + try: + return _validate_openapi_spec(json.loads(content)) + except json.JSONDecodeError: + try: + return _validate_openapi_spec(yaml.load(content, Loader=_OpenAPISafeLoader)) + except (yaml.YAMLError, RecursionError) as exc: + raise ValueError("Invalid YAML OpenAPI spec") from exc + + async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> _OpenAPISpec: if filepath.startswith("http://") or filepath.startswith("https://"): client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) @@ -208,14 +298,17 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None else await async_safe_get(client, filepath, max_response_bytes=max_bytes) ) r.raise_for_status() - return r.json() + return await asyncio.to_thread(_parse_openapi_spec, r.content) # fallback: local file - # Local filesystem path + return await asyncio.to_thread(_load_openapi_file, filepath) + + +def _load_openapi_file(filepath: str) -> _OpenAPISpec: if not os.path.exists(filepath): raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") with open(filepath, "r", encoding="utf-8") as f: - return json.load(f) + return _parse_openapi_spec(f.read()) def get_base_url(spec: _OpenAPISpec, spec_path: str | None = None) -> str: diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index c42579dc7f7..ac79c10869b 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -5405,7 +5405,42 @@ class TestMCPServerManager: ( httpx.Response(200, text="secret invalid JSON body"), "unhealthy", - "OpenAPI specification could not be loaded (JSONDecodeError)", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="paths: [secret invalid YAML body"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="paths: {}\nx: \x01"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: &a []\ny: *a"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!bool invalid"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!int ''"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!float ''"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!timestamp invalid"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", ), ], ) diff --git a/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py b/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py index aac85160eba..0f29a3bf20d 100644 --- a/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py +++ b/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py @@ -1,9 +1,15 @@ from __future__ import annotations -from typing import Any, Dict +import asyncio +import json +import threading +from pathlib import Path +from typing import Any, Dict, Final import httpx import pytest +import yaml +from respx import MockRouter from litellm.proxy.mcp_server import openapi_to_mcp_generator as gen @@ -98,3 +104,147 @@ def test_load_openapi_spec_supports_local_file_path( spec = gen.load_openapi_spec(str(p)) assert spec == expected + + +@pytest.mark.parametrize("source", ["url", "file"]) +@pytest.mark.parametrize("suffix", [".yaml", ".json", ""]) +@pytest.mark.parametrize("document_format", ["json", "yaml"]) +@pytest.mark.parametrize("use_async", [False, True]) +def test_openapi_formats_preserve_document_across_sources( + tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + source: str, suffix: str, document_format: str, use_async: bool, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + expected: Final = { + "openapi": "3.0.0", + "paths": {"/ping": {"get": { + "operationId": "on", + "parameters": [ + {"name": "on", "in": "query", "schema": {"type": "string"}}, + {"name": "off", "in": "query", "schema": {"type": "string"}}, + ], + }}}, + "x-values": [100.0, "on", "off", "yes", "no", "2026-01-01", "1:20", True, False, None, 12, 10, 10], + } + content: Final = ( + json.dumps(expected).replace("100.0", "1e2") + if document_format == "json" + else "openapi: 3.0.0\npaths:\n /ping:\n get:\n operationId: on\n" + " parameters:\n" + " - {name: on, in: query, schema: {type: string}}\n" + " - {name: off, in: query, schema: {type: string}}\n" + "x-values: [1e2, on, off, yes, no, 2026-01-01, 1:20, TRUE, False, null, 012, 0o12, 0xA]\n" + ) + path: Final = tmp_path / ("spec" + suffix) + path.write_text(content, encoding="utf-8") + target: Final = str(path) if source == "file" else "https://93.184.216.34/spec" + suffix + if source == "url": + respx_mock.get(target).respond(200, text=content, headers={"content-type": "text/plain"}) + + spec: Final = asyncio.run(gen.load_openapi_spec_async(target)) if use_async else gen.load_openapi_spec(target) + + assert spec == expected + operation: Final = spec["paths"]["/ping"]["get"] + assert gen.sanitize_openapi_tool_name(operation["operationId"]) == "on" + assert gen.build_input_schema(operation) == { + "type": "object", + "properties": {"on": {"type": "string", "description": ""}, "off": {"type": "string", "description": ""}}, + "required": [], + } + + +_LONG_YAML_INTEGER: Final = "1" * 1025 + + +@pytest.mark.parametrize("content,message", [ + ("", "OpenAPI spec must be a JSON or YAML object"), + ("null", "OpenAPI spec must be a JSON or YAML object"), + ("[]", "OpenAPI spec must be a JSON or YAML object"), + ("- item", "OpenAPI spec must be a JSON or YAML object"), + ("plain text", "OpenAPI spec must be a JSON or YAML object"), + ("paths: [", "Invalid YAML OpenAPI spec"), + ("paths: {}\nx: \x01", "Invalid YAML OpenAPI spec"), + ("paths: " + "[" * 600 + "0" + "]" * 600, "Invalid YAML OpenAPI spec"), + ("x: !!python/object/apply:builtins.str [unsafe]", "Invalid YAML OpenAPI spec"), + ("x: &a []\ny: *a", "YAML aliases are not supported in OpenAPI specs"), + ("x: " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"), + ("x: !!int " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"), + ("x: !!int {=: '" + _LONG_YAML_INTEGER + "'}", "Invalid YAML OpenAPI spec"), + ("x: !!int [1]", "Invalid YAML OpenAPI spec"), + ("x: !!bool invalid", "Invalid YAML OpenAPI spec"), + ("x: !!bool []", "Invalid YAML OpenAPI spec"), + ("x: !!str []", "Invalid YAML OpenAPI spec"), + ("x: !!int ''", "Invalid YAML OpenAPI spec"), + ("x: !!float ''", "Invalid YAML OpenAPI spec"), + ("x: !!float {=: 1.2}", "Invalid YAML OpenAPI spec"), + ("x: !!timestamp invalid", "Invalid YAML OpenAPI spec"), + ("x: !!timestamp 2026-01-01", "Invalid YAML OpenAPI spec"), + ("x: !!set {a: null}", "Invalid YAML OpenAPI spec"), + ("x: !!binary YQ==", "Invalid YAML OpenAPI spec"), + ("x: {1: value}", "Invalid YAML OpenAPI spec"), + ("x: {[key]: value}", "Invalid YAML OpenAPI spec"), + ("x: .inf", "Invalid YAML OpenAPI spec"), + ("x: .nan", "Invalid YAML OpenAPI spec"), + ("x: 1e999", "Invalid YAML OpenAPI spec"), +], ids=[ + "empty", "null", "json-array", "yaml-sequence", "scalar", "syntax", "reader", "nesting", + "python-tag", "alias", "implicit-int", "explicit-int", "mapping-int", "sequence-int", + "invalid-bool", "sequence-bool", "sequence-str", "empty-int", "empty-float", "mapping-float", "invalid-timestamp", + "timestamp", "set", "binary", "non-string-key", "complex-key", "infinity", "nan", "float-overflow", +]) +def test_openapi_loader_rejects_invalid_or_unsafe_documents( + tmp_path: Path, content: str, message: str, +) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text(content, encoding="utf-8") + + with pytest.raises(ValueError, match="^" + message + "$"): + gen.load_openapi_spec(str(path)) + + +def test_openapi_yaml_integer_boundary_preserves_quoted_strings(tmp_path: Path) -> None: + integer: Final = "1" * 1024 + path: Final = tmp_path / "spec.yaml" + path.write_text("x-integer: " + integer + "\nx-string: '" + integer + "1'", encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x-integer": int(integer), "x-string": integer + "1"} + + +def test_openapi_json_integer_preserved_even_with_yaml_extension(tmp_path: Path) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text('{"x-integer":' + _LONG_YAML_INTEGER + "}", encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x-integer": int(_LONG_YAML_INTEGER)} + + +@pytest.mark.parametrize("value", ["! " + _LONG_YAML_INTEGER, "! '" + _LONG_YAML_INTEGER + "'"], ids=["plain", "quoted"]) +def test_openapi_nonspecific_tag_disables_scalar_resolution(tmp_path: Path, value: str) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text("x: " + value, encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x": _LONG_YAML_INTEGER} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["url", "file"]) +async def test_openapi_yaml_parses_off_the_event_loop( + tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, source: str, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + event_loop_thread: Final = threading.get_ident() + parser_threads: Final[list[int]] = [] + original_load: Final = yaml.load + + def observed_load(stream: str | bytes, Loader: type[yaml.SafeLoader]) -> object: + parser_threads.append(threading.get_ident()) + return original_load(stream, Loader=Loader) + + monkeypatch.setattr(yaml, "load", observed_load) + path: Final = tmp_path / "spec.yaml" + path.write_text("paths: {}", encoding="utf-8") + target: Final = str(path) if source == "file" else "https://93.184.216.34/spec.yaml" + if source == "url": + respx_mock.get(target).respond(200, text="paths: {}") + + assert await gen.load_openapi_spec_async(target) == {"paths": {}} + assert len(parser_threads) == 1 and parser_threads[0] != event_loop_thread From 2a9c63018242f014b5a7416a1e7b065144d10b1e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:47:04 -0400 Subject: [PATCH 023/147] fix(proxy): release concurrency slots when synthetic streams complete (#48628) --- litellm/proxy/common_request_processing.py | 7 +- litellm/proxy/utils.py | 21 +--- .../agent_endpoints/test_a2a_endpoints.py | 3 +- .../test_streaming_model_restamp.py | 53 ++++---- .../hooks/test_parallel_request_limiter_v3.py | 18 ++- tests/unit/proxy/test_budget_reservation.py | 3 +- .../proxy/test_common_request_processing.py | 117 ++++++++++++++++++ 7 files changed, 168 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3b9c17baa02..740d8742e5a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4355,6 +4355,8 @@ class ProxyBaseLLMRequestProcessing: proxy_logging_obj: ProxyLogging | None = None, ) -> None: with anyio.CancelScope(shield=True): + if stream_completed and proxy_logging_obj is not None and user_api_key_dict is not None: + await proxy_logging_obj.arelease_max_parallel_requests(user_api_key_dict) should_record_client_disconnect: Final = client_disconnected or (not stream_completed) recorded_client_disconnect = False if should_record_client_disconnect: @@ -4369,10 +4371,7 @@ class ProxyBaseLLMRequestProcessing: # A disconnect-time success event (the deferred-guardrail flush # above, or the partial-spend billing below) releases the # request's max_parallel_requests slot through the limiter's - # own success callback. Release the slot explicitly only when - # no such event fires, so exactly one release happens; two - # concurrent releases would race and double-decrement under the - # limiter's in-memory fallback. + # own success callback. Release explicitly when no event fires. success_event_owns_slot_release = deferred_stream_logging_armed if not deferred_stream_logging_armed: success_event_owns_slot_release = await _bill_partial_streamed_spend_on_disconnect( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d3202368d76..452434e8c81 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4942,31 +4942,22 @@ class ProxyLogging: logging_obj._deferred_stream_complete_args = None return True - async def arelease_max_parallel_requests_on_disconnect( + async def arelease_max_parallel_requests( self, user_api_key_dict: UserAPIKeyAuth, ) -> None: """ - Release the api-key max_parallel_requests slot when a streaming - response is cancelled mid-flight (client disconnect) and no logging - callback fired for it. Neither the success nor failure callback runs on - the resulting CancelledError / GeneratorExit, so the pre-call +1 would - otherwise leak. - - Awaited from the shielded streaming cleanup rather than scheduled - fire-and-forget, so the caller can make it the single owner of the - release: when a disconnect-time success event does fire (partial-spend - billing or a deferred-guardrail flush), that event's own limiter - callback releases the slot and this is not called at all. Two - concurrent releases of the same acquisition would otherwise race and - double-decrement under the limiter's in-memory fallback. + Release an ended request's v3 concurrency acquisition independently of logging. + The request stash serializes release with callbacks and clears the acquisition, + so repeated cleanup cannot release another request's slot. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) - _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests_on_disconnect + arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests + _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests async def arelease_responses_websocket_slot(self, user_api_key_dict: UserAPIKeyAuth) -> None: """Release the slot a Responses WebSocket took at admission. Only call this once that admission succeeded: diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index a2448694dd7..a7d7ac5627c 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -1058,6 +1058,7 @@ async def test_subscribe_to_task_calls_pre_call_hook(): """tasks/resubscribe must run pre_call_hook so guardrails configured on the agent are enforced before streaming begins.""" from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging agent = _make_agent_mock() mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) @@ -1083,7 +1084,7 @@ async def test_subscribe_to_task_calls_pre_call_hook(): async for chunk in response: yield chunk - mock_proxy_logging = MagicMock() + mock_proxy_logging: Final = MagicMock(spec=ProxyLogging) mock_proxy_logging.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data ) diff --git a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py index b7bc670c7f8..85bb8304b95 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -3,15 +3,18 @@ Tests for restamping the public model on Anthropic Messages streaming chunks. """ import json -from unittest.mock import AsyncMock, MagicMock +from collections.abc import AsyncIterator, Sequence import pytest +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, restamp_anthropic_stream_chunk_model, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging def _message_start_frame(model: str, line_end: str = "\n") -> bytes: @@ -22,15 +25,9 @@ def _message_start_frame(model: str, line_end: str = "\n") -> bytes: return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode() -def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: - async def _iterator_hook(**_kwargs): - for frame in frames: - yield frame - - proxy_logging_obj = MagicMock() - proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook - proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) - return proxy_logging_obj +async def _stream_frames(frames: Sequence[bytes]) -> AsyncIterator[bytes]: + for frame in frames: + yield frame def _model_from_frame(frame: bytes | str) -> str: @@ -76,15 +73,14 @@ def test_leaves_chunks_without_a_model_untouched(chunk): async def test_sse_generator_publishes_requested_model_on_message_start(): """The message_start event reports the requested model, not the provider's.""" delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' - proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((_message_start_frame("claude-haiku-4-5-20251001"), delta_frame)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] @@ -95,15 +91,13 @@ async def test_sse_generator_publishes_requested_model_on_message_start(): @pytest.mark.asyncio async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): - proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")]) - chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((_message_start_frame("claude-haiku-4-5-20251001"),)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), ) ] @@ -175,15 +169,14 @@ def test_dict_message_start_disarms_after_restamp(): @pytest.mark.asyncio async def test_sse_generator_restamps_message_start_split_across_chunks(): frame = _message_start_frame("claude-haiku-4-5-20251001") - proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((frame[:30], frame[30:])), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] @@ -250,15 +243,14 @@ def test_flush_disarms_the_restamper(): @pytest.mark.asyncio async def test_sse_generator_flushes_held_tail_at_end_of_stream(): unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] - proxy_logging_obj = _proxy_logging_obj_streaming([unterminated]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((unterminated,)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] @@ -271,15 +263,14 @@ async def test_sse_generator_flushes_held_tail_at_end_of_stream(): async def test_sse_generator_restamps_crlf_stream(): frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' - proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((frame, delta)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index b8dba5e1e68..cd0d786f841 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4378,9 +4378,12 @@ async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_hea @pytest.mark.asyncio -async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: +@pytest.mark.parametrize("release_from", ["post_call", "stream_cleanup"]) +async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails(release_from: str) -> None: from unittest.mock import AsyncMock + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + api_key = hash_token("sk-concurrent-release") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( @@ -4409,9 +4412,20 @@ async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redi counter_keys=[parallel_key], ) data = {"litellm_call_id": call_id} + logger: Final = ProxyLogging(user_api_key_cache=local_cache) + logger.proxy_hook_mapping["parallel_request_limiter"] = handler post_call_task = asyncio.create_task( - handler.async_post_call_success_hook( + ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, + request_data=data, + response=None, + stream_completed=True, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=logger, + ) + if release_from == "stream_cleanup" + else handler.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=ResponsesAPIResponse( diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 42b41a2e3f2..535c3da8c3a 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -4,6 +4,7 @@ import time from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -3233,7 +3234,7 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat async def one_chunk(user_api_key_dict, response, request_data): yield {"content": "hi"} - streaming_logging_obj = MagicMock() + streaming_logging_obj: Final = MagicMock(spec=ProxyLogging) streaming_logging_obj.async_post_call_streaming_iterator_hook = one_chunk streaming_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs["response"] diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5c99e5a01ea..52816dc50a8 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1,4 +1,5 @@ import asyncio +import contextvars import copy import datetime import json @@ -12055,6 +12056,122 @@ class TestResponseCostFromLoggingObjForCacheHits: assert cost == "" +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["messages", "openai"]) +@pytest.mark.parametrize("keepalive", ["disabled", "fast", "pinged"]) +async def test_completed_synthetic_stream_reuses_capacity_without_freeing_another_request( + surface: Literal["messages", "openai"], + keepalive: Literal["disabled", "fast", "pinged"], + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache, hash_token + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse + + def clock() -> datetime.datetime: + return datetime.datetime(2026, 1, 1) + + cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(cache), time_provider=clock + ) + logger: Final = ProxyLogging(user_api_key_cache=cache) + logger.proxy_hook_mapping["parallel_request_limiter"] = limiter + auth: Final = ProxyUserAPIKeyAuth(api_key=hash_token("sk-synthetic-stream"), max_parallel_requests=2) + counter: Final = f"{{api_key:{auth.api_key}}}:max_parallel_requests" + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logger) + + async def admit(call_id: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": "test-model", "litellm_call_id": call_id}, + call_type="anthropic_messages", + ) + + async def request(call_id: str) -> None: + upstream_answers: Final = asyncio.Event() + if keepalive != "pinged": + upstream_answers.set() + + async def produce() -> StreamingResponse: + await admit(call_id) + with pytest.raises(HTTPException) as rejected: + await contextvars.Context().run(asyncio.create_task, admit("overlapping")) + assert rejected.value.status_code == 429 + await upstream_answers.wait() + body: Final = AnthropicMessagesResponse( + id=call_id, + type="message", + role="assistant", + model="test-model", + content=[{"type": "text", "text": "Search completed"}], + stop_reason="end_turn", + stop_sequence=None, + usage={"input_tokens": 0, "output_tokens": 0}, + ) + source: Final = FakeAnthropicMessagesStreamIterator(body) + data: Final = {"model": "test-model", "litellm_call_id": call_id} + stream: Final = ( + ProxyBaseLLMRequestProcessing.async_sse_data_generator(source, auth, data, logger) + if surface == "messages" + else proxy_server.async_data_generator(source, auth, data) + ) + return StreamingResponse(stream) + + response: Final = await open_sse_before_first_byte( + produce(), None if keepalive == "disabled" else 0.001 if keepalive == "pinged" else 60.0 + ) + first: Final = await anext(response.body_iterator) + if keepalive == "pinged": + assert first == SSE_COMMENT_PING_BYTES + upstream_answers.set() + frames: Final = (first, *tuple([frame async for frame in response.body_iterator])) + wire: Final = b"".join(frame.encode() if isinstance(frame, str) else frame for frame in frames) + assert b'"type": "message_stop"' in wire + + await contextvars.Context().run(asyncio.create_task, admit("still-streaming")) + for index in range(4): + await contextvars.Context().run(asyncio.create_task, request(f"completed-{index}")) + occupied: Final = await cache.async_get_cache(counter) + assert isinstance(occupied, dict) and len(occupied) == 1, occupied + + +@pytest.mark.asyncio +async def test_completed_stream_cleanup_does_not_release_an_unowned_legacy_slot() -> None: + from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler + from litellm.proxy.utils import InternalUsageCache, hash_token + + def clock() -> datetime.datetime: + return datetime.datetime(2026, 1, 1) + + cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), clock=clock) + logger: Final = ProxyLogging(user_api_key_cache=cache) + logger.proxy_hook_mapping["parallel_request_limiter"] = limiter + auth: Final = ProxyUserAPIKeyAuth(api_key=hash_token("sk-legacy-stream"), max_parallel_requests=1) + await limiter.async_pre_call_hook(auth, cache, {}, "acompletion") + + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, + request_data={}, + response=None, + stream_completed=True, + user_api_key_dict=auth, + proxy_logging_obj=logger, + ) + + with pytest.raises(HTTPException) as rejected: + await limiter.async_pre_call_hook(auth, cache, {}, "acompletion") + assert rejected.value.status_code == 429 + + def test_maybe_get_model_id_reads_the_chat_metadata_the_router_updated_in_place(): processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"metadata": {"model_info": {"id": "declared-text"}}}) From a53ea5bd7784c51e4c0b726efe55de44c56c8b5d Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 27 Oct 2026 14:49:07 -0700 Subject: [PATCH 024/147] fix(proxy): keep every startup line JSON when json_logs is on (#47035) * fix(proxy): keep every startup line JSON when json_logs is on Boot-time messages the proxy printed straight to stdout (banner, feedback box, model list, cache and callback notices, policy table, migration notices, CLI launch notes) and the first records of uvicorn- or granian-spawned workers came out as plain text while json_logs was on, breaking log pipelines that parse every line as JSON. - print_startup_message() prints unchanged in plain mode and, with json_logs on, emits one INFO record on the proxy logger (ANSI stripped) that skips the logger and handler level checks like the print it replaces; show_banner() routes to it in JSON mode only - the CLI exports JSON_LOGS=true when the config sets json_logs so spawned workers log JSON from their first record - the --enforce_prisma_migration_check deprecation notice prints after the config is read so a config-only json_logs covers it Resolves LIT-1624 * fix(proxy): keep show_banner() SDK-free, pick the JSON banner in initialize() Since #46803 the `lite` shell calls show_banner() against an uninitialized litellm package, and reading litellm.json_logs inside it initialized the whole SDK at the banner (0.14 s to 0.83 s for a bare `lite` here, the seconds that PR removed on small VMs). show_banner() is back to its main version; the proxy's initialize() now emits the JSON banner record itself and calls show_banner() only in plain mode. The banner unit tests move to the initialize() call site. * fix(proxy): keep the JSON banner when an operator --log_config names the proxy logger uvicorn applies --log_config after proxy_cli has switched JSON logging on and before initialize() switches it on again through load_config(). A logging config that names "LiteLLM Proxy" removes that logger's handlers in the window between the two, which is where initialize() emits the banner and the feedback box, so at one worker both went nowhere; two-worker children import litellm after uvicorn's logging setup and were not affected. print_startup_message() now falls back to a one-off JSON handler, on the stream the JSON handler uses, when the proxy logger has no handler of its own. Live, json_logs on plus a log_config naming the proxy logger, one worker: main prints the banner and the feedback box as plaintext, e2303a3925 prints neither, this head prints both as one JSON record each. New unit test fails on e2303a3925; new integration cell test_json_logs_with_operator_log_config_keeps_every_line_json passes at one and two workers. * fix(proxy): keep a long startup listing whole in its JSON record print_startup_message ran its record through the proxy logger's filters, so in JSON mode the required-migration SQL list and the policy list were middle-truncated at MAX_STRING_LENGTH_STDOUT_LOG (4096 by default) with the litellm_truncated marker, where the prints they replace showed them whole. That cap bounds echoed request payloads; a startup message is operator text, so the record now gets the uncapped processing of a DEBUG line: secret redaction and base64 collapse still apply, the text cap does not. Red on 41f92b745a, green here: - tests/unit/test_logging.py: a listing longer than the cap comes out whole; a credential in a startup message is still redacted - tests/integration/configuration/test_json_logs_startup.py: a 60-policy config boots with json_logs on and the policy listing record (9762 chars) carries every policy with no truncation marker (4096 chars, marker, on the unfixed tree) * fix(proxy): keep the --enforce_prisma_migration_check notice on --health and --test This PR moved the deprecation notice after the config is read, so that a json_logs set only in the config applies to it; --health and --test return before that point and stopped printing it, a visible change main did not have. The notice is now also emitted before those two returns, through the same helper, so a deployment running `litellm --health --enforce_prisma_migration_check` still learns the flag is dead, as on main. A normal server start emits it once, after the config, as in the previous commit. Red on 2d7a2d6844, green here: tests/unit/proxy/test_proxy_cli.py asserts the notice on --health and on --test. * fix(proxy): keep the JSON startup records on stdout under an operator handler on the proxy logger An operator --log_config that gives "LiteLLM Proxy" a handler of its own, a NullHandler that silences it or a plain WARNING-level handler on stderr, is in force in the window where the banner, the feedback box and the cache notice are emitted: uvicorn applies it after the CLI process switched JSON on and before initialize() re-installs the JSON handler through load_config(). 41f92b745a sent the startup records through the proxy logger's handlers and fell back to a JSON stdout handler only when there were none, so under such a config the records landed in the operator's handler, as plaintext on stderr, or vanished, where the prints they replace reached stdout whatever the logging config said. print_startup_message() now writes every startup record through the proxy's own JSON handler on stdout, filters included, and never through the proxy logger's handlers; the operator's handler gets nothing from a startup message, as it got nothing from the print. The normal JSON handler is the same class, formatter and filters, so the records a pipeline reads are unchanged when no operator config is in play. Tests: tests/unit/test_logging.py gets the two handler shapes (red on f6f5957d16: the record went to stderr as plaintext or nowhere); tests/integration/configuration/ test_json_logs_startup.py gets a silencing operator --log_config cell at 1 and 2 workers (red on f6f5957d16: banner, feedback box and cache notice missing; green here). * fix(proxy): print the --enforce_prisma_migration_check notice when the config fails to load f6f5957d16 printed the deprecation notice after the config is read, so a json_logs set only in the config makes it a JSON record. A boot whose config file is missing or does not parse then exited before that point and never printed it, where main printed it before the error (10-20 gauntlet finding). The notice is printed where main printed it, unless a config file can still switch json_logs on: then it is held until the config is read and printed right after the json_logs switch, and a click close callback prints it on the way out if the boot fails first. functools.cache keeps it to one print either way. Tests: tests/unit/proxy/test_proxy_cli.py covers a failed config read (red on 75043aed83), a plain config printing it once and a config-only json_logs making it one JSON record; tests/integration/configuration/test_json_logs_startup.py adds the JSON cell and a missing and an unparseable config cell (red on 75043aed83, green here), through refused_boot_log, which now takes extra CLI arguments. * fix(proxy): keep startup records as printed instead of running secret redaction In JSON mode the startup messages ran secret redaction, which the prints they replace never had. Its sk- key pattern has no left boundary, so a model named risk-assessment-gpt4o was logged as riREDACTED in the model list (10-21 gauntlet finding). The records now skip redaction like the prints did; the processed record still counts as scrubbed, so the handler's secret filter and the formatter leave it alone. The stdout cap stays off and base64 collapse stays on, as before. Tests: tests/unit/test_logging.py checks the model name comes out whole (red on 7d17f126e5), replacing the test that asserted startup records were redacted; tests/integration/configuration/test_json_logs_startup.py adds a boot with a risk-assessment- model whose model list record carries the name (red on 7d17f126e5). * refactor(proxy): meet the AGENTS.md rules Greptile flagged on the startup-message change The policy listing builds its rows as tuples in two helpers instead of appending to a list; plain output is unchanged (same bytes as before across 3000 random policy and attachment configs). The explanatory comments added with this change are gone, and the docstring of print_startup_message now says the record is whole and unredacted. Tests: the migration-diff JSON test no longer runs the fake prisma CLI; it checks the skipped-diff notice in-process, and the fixture extension it needed is reverted. The json_logs export and deprecation-notice CLI tests run the real proxy module and a real config file instead of a mocked ProxyConfig and proxy module, and the --health/--test notice tests no longer patch our helpers. Each still fails when the export, the hold or call_on_close is removed. New test parameters are typed. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/_logging.py | 44 ++ litellm/proxy/db/check_migration.py | 19 +- litellm/proxy/policy_engine/init_policies.py | 111 ++--- litellm/proxy/proxy_cli.py | 113 +++-- litellm/proxy/proxy_server.py | 70 +-- tests/integration/_support/process.py | 2 + .../configuration/test_json_logs_startup.py | 397 ++++++++++++++++++ tests/unit/proxy/db/test_check_migration.py | 36 ++ .../proxy/policy_engine/test_init_policies.py | 116 +++++ tests/unit/proxy/test_proxy_cli.py | 138 ++++++ ...test_proxy_server_endpoints_and_startup.py | 62 ++- tests/unit/test_logging.py | 220 ++++++++++ 12 files changed, 1188 insertions(+), 140 deletions(-) create mode 100644 tests/integration/configuration/test_json_logs_startup.py create mode 100644 tests/unit/proxy/policy_engine/test_init_policies.py diff --git a/litellm/_logging.py b/litellm/_logging.py index d08254db3fd..8974f85f517 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,6 +1,7 @@ import ast import contextvars import functools +import inspect import itertools import logging import os @@ -897,6 +898,49 @@ verbose_proxy_logger.addFilter(_diagnostic_filter) verbose_proxy_stdout_logger.addFilter(_diagnostic_filter) verbose_logger.addFilter(_diagnostic_filter) +_ANSI_ESCAPE: Final = re.compile(r"\x1b\[[0-9;]*m") + + +def print_startup_message(message: str) -> None: + """Write a boot-time message meant for the operator's terminal. + + Plain mode prints it unchanged, colors and spacing included. With JSON logs on, every line of + stdout is read by a log pipeline, so the message goes out as one INFO record of the proxy logger + instead, on stdout like the print whatever handlers the logging config gave that logger, with the + ANSI color codes and surrounding blank lines removed, located at the caller and emitted at any log + level, whole and unredacted as the print showed it; a blank spacer produces no record. + """ + if not litellm.json_logs: + print(message) # noqa: T201 # plain mode: this is the operator-terminal print it replaces + return + text: Final = _ANSI_ESCAPE.sub("", message).strip("\n") + if not text: + return + frame: Final = inspect.currentframe() + caller: Final = frame.f_back if frame is not None else None + record: Final = verbose_proxy_logger.makeRecord( + verbose_proxy_logger.name, + logging.INFO, + caller.f_code.co_filename if caller is not None else __file__, + caller.f_lineno if caller is not None else 0, + text, + (), + None, + func=caller.f_code.co_name if caller is not None else None, + ) + if not _process_record(record, base64_limit=_get_max_base64_length_stdout_log(), text_limit=0, redact=False): + return + _startup_json_handler().handle(record) + + +def _startup_json_handler() -> logging.Handler: + """The proxy's JSON handler on stdout, filters included, for one startup message.""" + handler: Final = LevelRoutingStreamHandler() + handler.setFormatter(JsonFormatter()) + handler.addFilter(_secret_filter) + handler.addFilter(_correlation_filter) + return handler + def _suppress_loggers(): """Suppress noisy loggers at INFO level""" diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index 6e2a06c96e1..83cc4501466 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -4,7 +4,7 @@ import os import subprocess from typing import Final -from litellm._logging import verbose_logger +from litellm._logging import print_startup_message, verbose_logger def extract_sql_commands(diff_output: str) -> list[str]: @@ -63,8 +63,8 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: run_prisma, ) except ImportError as e: - print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages - f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" + print_startup_message( + f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}", ) return False, [] @@ -89,22 +89,21 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: sql_commands: Final = extract_sql_commands(result.stdout) if sql_commands: - print("Changes to DB Schema detected") # noqa: T201 - print("Required SQL commands:") # noqa: T201 - for command in sql_commands: - print(command) # noqa: T201 + print_startup_message( + "\n".join(("Changes to DB Schema detected", "Required SQL commands:", *sql_commands)), + ) return True, sql_commands else: return False, [] except subprocess.TimeoutExpired: - print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + print_startup_message( f"Timed out after {timeout}s generating the migration diff. " - f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer." + f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer.", ) return False, [] except subprocess.CalledProcessError as e: error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" - print(error_message) # noqa: T201 + print_startup_message(error_message) return False, [] diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index d75e651e6c7..59d15b0ba0a 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -6,10 +6,12 @@ Configuration structure: - policy_attachments: Define WHERE policies apply (teams, keys, models) """ +import sys from collections.abc import Mapping, Sequence +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Optional -from litellm._logging import verbose_proxy_logger +from litellm._logging import print_startup_message, verbose_proxy_logger from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_validator import PolicyValidator @@ -25,6 +27,42 @@ _yellow_color_code: Final = "\033[93m" _reset_color_code: Final = "\033[0m" +def _policy_lines(policy_name: str, policy_data: Mapping[str, object]) -> tuple[str, ...]: + guardrails: Final = policy_data.get("guardrails", {}) + inherit: Final = policy_data.get("inherit") + condition: Final = policy_data.get("condition") + description: Final = policy_data.get("description") + + guardrails_add: Final = guardrails.get("add", []) if isinstance(guardrails, dict) else [] + guardrails_remove: Final = guardrails.get("remove", []) if isinstance(guardrails, dict) else [] + model_condition: Final = condition.get("model") if isinstance(condition, dict) else None + inherit_str: Final = f" (inherits: {inherit})" if inherit else "" + + return ( + f"{_blue_color_code} - {policy_name}{inherit_str}{_reset_color_code}", + *((f" description: {description}",) if description else ()), + *((f" guardrails.add: {guardrails_add}",) if guardrails_add else ()), + *((f" guardrails.remove: {guardrails_remove}",) if guardrails_remove else ()), + *((f" condition.model: {model_condition}",) if model_condition else ()), + ) + + +def _attachment_line(attachment: Mapping[str, object]) -> str: + policy: Final = attachment.get("policy", "unknown") + scope: Final = attachment.get("scope") + teams: Final = attachment.get("teams") + keys: Final = attachment.get("keys") + models: Final = attachment.get("models") + + scope_parts: Final = ( + *(("scope=* (global)",) if scope == "*" else ()), + *((f"teams={teams}",) if teams else ()), + *((f"keys={keys}",) if keys else ()), + *((f"models={models}",) if models else ()), + ) + return f" - {policy} -> {', '.join(scope_parts) if scope_parts else 'all'}" + + def _print_policies_on_startup( policies_config: Mapping[str, Mapping[str, object]], policy_attachments_config: Sequence[Mapping[str, object]] | None = None, @@ -32,67 +70,32 @@ def _print_policies_on_startup( """ Print loaded policies to console on startup (similar to model list). """ - import sys - - print( # noqa: T201 + header: Final = ( f"{_green_color_code}\nLiteLLM Policy Engine: Loaded {len(policies_config)} policies{_reset_color_code}\n" ) - sys.stdout.flush() - - for policy_name, policy_data in policies_config.items(): - guardrails = policy_data.get("guardrails", {}) - inherit = policy_data.get("inherit") - condition = policy_data.get("condition") - description = policy_data.get("description") - - guardrails_add = guardrails.get("add", []) if isinstance(guardrails, dict) else [] - guardrails_remove = guardrails.get("remove", []) if isinstance(guardrails, dict) else [] - inherit_str = f" (inherits: {inherit})" if inherit else "" - - print( # noqa: T201 - f"{_blue_color_code} - {policy_name}{inherit_str}{_reset_color_code}" + policy_lines: Final = tuple( + chain.from_iterable( + _policy_lines(policy_name, policy_data) for policy_name, policy_data in policies_config.items() ) - if description: - print(f" description: {description}") # noqa: T201 - if guardrails_add: - print(f" guardrails.add: {guardrails_add}") # noqa: T201 - if guardrails_remove: - print(f" guardrails.remove: {guardrails_remove}") # noqa: T201 - if condition: - model_condition = condition.get("model") if isinstance(condition, dict) else None - if model_condition: - print(f" condition.model: {model_condition}") # noqa: T201 + ) - # Print attachments if policy_attachments_config: - print( # noqa: T201 - f"\n{_yellow_color_code}Policy Attachments: {len(policy_attachments_config)} attachment(s){_reset_color_code}" + print_startup_message( + "\n".join( + ( + header, + *policy_lines, + f"\n{_yellow_color_code}Policy Attachments: {len(policy_attachments_config)} attachment(s){_reset_color_code}", + *(_attachment_line(attachment) for attachment in policy_attachments_config), + "", + ) + ) ) - for attachment in policy_attachments_config: - policy = attachment.get("policy", "unknown") - scope = attachment.get("scope") - teams = attachment.get("teams") - keys = attachment.get("keys") - models = attachment.get("models") - - scope_parts = [] - if scope == "*": - scope_parts.append("scope=* (global)") - if teams: - scope_parts.append(f"teams={teams}") - if keys: - scope_parts.append(f"keys={keys}") - if models: - scope_parts.append(f"models={models}") - scope_str = ", ".join(scope_parts) if scope_parts else "all" - - print(f" - {policy} -> {scope_str}") # noqa: T201 else: - print( # noqa: T201 - f"\n{_yellow_color_code}Warning: No policy_attachments configured. Policies will not be applied to any requests.{_reset_color_code}" + print_startup_message("\n".join((header, *policy_lines))) + print_startup_message( + f"\n{_yellow_color_code}Warning: No policy_attachments configured. Policies will not be applied to any requests.{_reset_color_code}\n", ) - - print() # noqa: T201 sys.stdout.flush() diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index fc33ab1bd0b..909f6de40cd 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1,4 +1,5 @@ # ruff: noqa: T201 +import functools import importlib import json import os @@ -7,7 +8,7 @@ import re import subprocess import sys import urllib.parse as urlparse -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from pathlib import Path from typing import TYPE_CHECKING, Any, Final @@ -18,6 +19,7 @@ from dotenv import load_dotenv from pydantic import ConfigDict import litellm +from litellm._logging import print_startup_message from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY from litellm.proxy.db.pgbouncer import ( PgBouncerError, @@ -202,6 +204,29 @@ def deprecated_v2_flag_passed_on_cli() -> bool: return ctx.get_parameter_source("use_v2_migration_resolver") is ParameterSource.COMMANDLINE +def _print_enforce_prisma_migration_check_notice() -> None: + print_startup_message( + "\033[1;33mLiteLLM Proxy: --enforce_prisma_migration_check is " + "deprecated and has no effect, because the proxy always exits " + "when database setup fails at startup. You can safely remove it.\033[0m", + ) + + +def _print_or_hold_enforce_prisma_migration_check_notice( + *, reads_config_file: bool, returns_before_config: bool +) -> Callable[[], None] | None: + """Print the notice now, or, while a config file can still switch json_logs on, return a printer + to call once it is read; a boot that fails before that prints the held notice on the way out.""" + if not reads_config_file or litellm.json_logs or returns_before_config: + _print_enforce_prisma_migration_check_notice() + return None + print_once: Final = functools.cache(_print_enforce_prisma_migration_check_notice) + ctx: Final = click.get_current_context(silent=True) + if ctx is not None: + ctx.call_on_close(print_once) + return print_once + + class ProxyInitializationHelpers: @staticmethod def echo_litellm_version(): @@ -335,7 +360,7 @@ class ProxyInitializationHelpers: "timeout_graceful_shutdown": GracefulShutdownManager.get_timeout(), } if log_config is not None: - print(f"Using log_config: {log_config}") + print_startup_message(f"Using log_config: {log_config}") uvicorn_args["log_config"] = log_config elif litellm.json_logs: # Use JSON log config for uvicorn to ensure all logs (including exceptions) are JSON @@ -348,10 +373,10 @@ class ProxyInitializationHelpers: if "timeout_worker_healthcheck" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["timeout_worker_healthcheck"] = timeout_worker_healthcheck else: - print( + print_startup_message( f"\033[1;33mLiteLLM Proxy: --timeout_worker_healthcheck " f"requires uvicorn>=0.37.0, but installed uvicorn=={uvicorn.__version__}. " - f"Ignoring the flag.\033[0m" + f"Ignoring the flag.\033[0m", ) return uvicorn_args @@ -371,18 +396,18 @@ class ProxyInitializationHelpers: import uvicorn if max_requests_before_restart is None: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " - "has no effect without --max_requests_before_restart\033[0m\n" + "has no effect without --max_requests_before_restart\033[0m\n", ) return if "limit_max_requests_jitter" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["limit_max_requests_jitter"] = jitter else: - print( + print_startup_message( f"\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " f"requires uvicorn>=0.41.0, but installed uvicorn=={uvicorn.__version__}. " - f"Ignoring the flag.\033[0m" + f"Ignoring the flag.\033[0m", ) _apply_uvicorn_max_requests_jitter = apply_uvicorn_max_requests_jitter @@ -506,12 +531,12 @@ class ProxyInitializationHelpers: from hypercorn.asyncio import serve from hypercorn.config import Config - print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Hypercorn\033[0m\n") + print_startup_message(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Hypercorn\033[0m\n") config: Final = Config() config.bind = [f"{host}:{port}"] if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) config.certfile = ssl_certfile_path @@ -544,14 +569,14 @@ class ProxyInitializationHelpers: from granian import Granian from granian.constants import Interfaces - print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n") + print_startup_message(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n") if max_requests_before_restart is not None: - print( + print_startup_message( "\033[1;33mLiteLLM: --max_requests_before_restart is not supported by Granian " - "(Granian uses workers_lifetime in seconds, not a per-request limit).\033[0m\n" + "(Granian uses workers_lifetime in seconds, not a per-request limit).\033[0m\n", ) if ciphers is not None: - print("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") + print_startup_message("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") kwargs: Final[dict[str, object]] = { "target": "litellm.proxy.proxy_server:app", @@ -564,7 +589,7 @@ class ProxyInitializationHelpers: if granian_runtime_threads is not None: kwargs["runtime_threads"] = granian_runtime_threads if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) kwargs["ssl_cert"] = Path(ssl_certfile_path) @@ -625,13 +650,15 @@ class ProxyInitializationHelpers: \n """ ) - print() - print( + print_startup_message("") + print_startup_message( '\033[1;34mLiteLLM: Test your local proxy with: "litellm --test" This runs an openai.ChatCompletion request to your proxy [In a new terminal tab]\033[0m\n' ) - print(f"\033[1;34mLiteLLM: Curl Command Test for your local proxy\n {curl_command} \033[0m\n") - print("\033[1;34mDocs: https://docs.litellm.ai/docs/simple_proxy\033[0m\n") - print(f"\033[1;34mSee all Router/Swagger docs on http://0.0.0.0:{port} \033[0m\n") + print_startup_message( + f"\033[1;34mLiteLLM: Curl Command Test for your local proxy\n {curl_command} \033[0m\n" + ) + print_startup_message("\033[1;34mDocs: https://docs.litellm.ai/docs/simple_proxy\033[0m\n") + print_startup_message(f"\033[1;34mSee all Router/Swagger docs on http://0.0.0.0:{port} \033[0m\n") def load_config(self): # note: This Loads the gunicorn config - has nothing to do with LiteLLM Proxy config @@ -651,7 +678,9 @@ class ProxyInitializationHelpers: # gunicorn app function return self.application - print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n") + print_startup_message( + f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n" + ) graceful_shutdown_timeout: Final = GracefulShutdownManager.get_timeout() gunicorn_options: Final = { "bind": f"{host}:{port}", @@ -669,9 +698,9 @@ class ProxyInitializationHelpers: gunicorn_options["max_requests"] = max_requests_before_restart if max_requests_before_restart_jitter is not None: if max_requests_before_restart is None: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " - "has no effect without --max_requests_before_restart\033[0m\n" + "has no effect without --max_requests_before_restart\033[0m\n", ) else: gunicorn_options["max_requests_jitter"] = max_requests_before_restart_jitter @@ -686,7 +715,7 @@ class ProxyInitializationHelpers: gunicorn_options["child_exit"] = child_exit if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) gunicorn_options["certfile"] = ssl_certfile_path @@ -710,7 +739,7 @@ class ProxyInitializationHelpers: with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - print(f""" + print_startup_message(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` """) @@ -779,7 +808,7 @@ class ProxyInitializationHelpers: os.makedirs(multiproc_dir, exist_ok=True) wipe_directory(multiproc_dir) action: Final = "Using existing" if configured_dir else "Auto-created" - print(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}") + print_startup_message(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}") return multiproc_dir _maybe_setup_prometheus_multiproc_dir = maybe_setup_prometheus_multiproc_dir @@ -1199,12 +1228,13 @@ def run_server( 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 " - "deprecated and has no effect, because the proxy always exits " - "when database setup fails at startup. You can safely remove it.\033[0m" + held_prisma_notice: Final = ( + _print_or_hold_enforce_prisma_migration_check_notice( + reads_config_file=config is not None, returns_before_config=health is True or test is True ) + if enforce_prisma_migration_check + else None + ) if model and "ollama" in model and api_base is None: ProxyInitializationHelpers.run_ollama_serve() if health is True: @@ -1320,8 +1350,11 @@ def run_server( import litellm litellm.json_logs = True + os.environ["JSON_LOGS"] = "true" litellm.turn_on_json() + if held_prisma_notice is not None: + held_prisma_notice() ### GENERAL SETTINGS ### general_settings = _config.get("general_settings", {}) if general_settings is None: @@ -1488,7 +1521,7 @@ def run_server( lifetime_params, ) ) - print( + print_startup_message( postgres_connection_budget_message( writer_limit=_url_query_value(os.getenv("DATABASE_URL"), "connection_limit") or f"{db_connection_pool_limit}", @@ -1520,19 +1553,19 @@ def run_server( env_value=os.getenv("USE_V2_MIGRATION_RESOLVER"), ) if deprecated_v2_flag_passed_on_cli() and use_v2_resolver: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: --use_v2_migration_resolver is " "deprecated and has no effect, because the v2 migration " "resolver is now the default. You can safely remove it. To " "opt back into the legacy v1 resolver, pass " - "--use_legacy_migration_resolver.\033[0m" + "--use_legacy_migration_resolver.\033[0m", ) if not use_v2_resolver: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration " "resolver. It performs the diff-and-force recovery that can " "cause schema thrashing during rolling deploys where two " - "LiteLLM versions contend for the same DB.\033[0m" + "LiteLLM versions contend for the same DB.\033[0m", ) try: migrated: Final = PrismaManager.setup_database( @@ -1600,7 +1633,7 @@ def run_server( # Skip server startup if requested (after all setup is done) if skip_server_startup: - print("LiteLLM: Setup complete. Skipping server startup as requested.") + print_startup_message("LiteLLM: Setup complete. Skipping server startup as requested.") return # Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups @@ -1619,7 +1652,7 @@ def run_server( ) except MetricsServerStartupError as error: raise click.ClickException(str(error)) from error - print( + print_startup_message( f"\033[1;32mLiteLLM: Serving Prometheus metrics on {host}:{prometheus_metrics_port}/metrics " f"(pid {metrics_process.pid})\033[0m" ) @@ -1642,7 +1675,7 @@ def run_server( run_gunicorn=run_gunicorn is True, ) if single_worker_warning is not None: - print(f"\033[1;33m{single_worker_warning}\033[0m\n") + print_startup_message(f"\033[1;33m{single_worker_warning}\033[0m\n") if run_gunicorn is False and run_hypercorn is False and run_granian is False: if limit_concurrency is not None: uvicorn_args["limit_concurrency"] = limit_concurrency @@ -1653,7 +1686,7 @@ def run_server( jitter=max_requests_before_restart_jitter, ) if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) uvicorn_args["ssl_keyfile"] = ssl_keyfile_path diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc79becf7bb..3b36ea5afce 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -273,25 +273,27 @@ def generate_feedback_box(): # Select a random message message: Final = random.choice(list_of_messages) - print() # noqa: T201 - print("\033[1;37m" + "#" + "-" * box_width + "#\033[0m") # noqa: T201 - print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa: T201 - print("\033[1;37m" + f"# {message:^59} #\033[0m") # noqa: T201 - print( # noqa: T201 - "\033[1;37m" + "# {:^59} #\033[0m".format("https://github.com/BerriAI/litellm/issues/new") + print_startup_message( + "\n".join( + ( + "", + "\033[1;37m" + "#" + "-" * box_width + "#\033[0m", + "\033[1;37m" + "#" + " " * box_width + "#\033[0m", + "\033[1;37m" + f"# {message:^59} #\033[0m", + "\033[1;37m" + "# {:^59} #\033[0m".format("https://github.com/BerriAI/litellm/issues/new"), + "\033[1;37m" + "#" + " " * box_width + "#\033[0m", + "\033[1;37m" + "#" + "-" * box_width + "#\033[0m", + "", + " Thank you for using LiteLLM! - Krrish & Ishaan", + "", + "", + "", + "\033[1;31mGive Feedback / Get Help: https://github.com/BerriAI/litellm/issues/new\033[0m", + "", + "", + ) + ) ) - print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa: T201 - print("\033[1;37m" + "#" + "-" * box_width + "#\033[0m") # noqa: T201 - print() # noqa: T201 - print(" Thank you for using LiteLLM! - Krrish & Ishaan") # noqa: T201 - print() # noqa: T201 - print() # noqa: T201 - print() # noqa: T201 - print( # noqa: T201 - "\033[1;31mGive Feedback / Get Help: https://github.com/BerriAI/litellm/issues/new\033[0m" - ) - print() # noqa: T201 - print() # noqa: T201 import contextlib @@ -304,6 +306,7 @@ import litellm._redis from litellm import Router from litellm._internal_context import service_target, with_service_target from litellm._logging import ( + print_startup_message, redact_string, verbose_proxy_logger, verbose_router_logger, @@ -6412,17 +6415,13 @@ class ProxyConfig: search_tools_parsed: Final[list[SearchToolTypedDict]] = [] - print( # noqa: T201 - "\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m" - ) + print_startup_message("\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m") for search_tool in search_tools_raw: # Display loaded search tool search_tool_name = search_tool.get("search_tool_name", "") search_provider = search_tool.get("litellm_params", {}).get("search_provider", "") - print( # noqa: T201 - f"\033[32m {search_tool_name} ({search_provider})\033[0m" - ) + print_startup_message(f"\033[32m {search_tool_name} ({search_provider})\033[0m") # Handle os.environ/ variables in litellm_params litellm_params = search_tool.get("litellm_params", {}) @@ -6598,7 +6597,7 @@ class ProxyConfig: reset_color_code: Final = "\033[0m" for key, value in litellm_settings.items(): if key == "cache" and value is True: - print(f"{blue_color_code}\nSetting Cache on Proxy") # noqa: T201 + print_startup_message(f"{blue_color_code}\nSetting Cache on Proxy") from litellm.caching.caching import Cache cache_params = {} @@ -6789,7 +6788,7 @@ class ProxyConfig: if PrometheusLogger is not None: verbose_proxy_logger.debug("mounting metrics endpoint") PrometheusLogger._mount_metrics_endpoint() - print( # noqa: T201 + print_startup_message( f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}" ) elif key == "failure_callback": @@ -6808,7 +6807,7 @@ class ProxyConfig: # these are litellm callbacks - "langfuse", "sentry", "wandb" else: litellm.logging_callback_manager.add_litellm_failure_callback(callback) - print( # noqa: T201 + print_startup_message( f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}" ) elif key == "audit_log_callbacks": @@ -6833,7 +6832,7 @@ class ProxyConfig: _store_audit_logs = litellm_settings.get("store_audit_logs", litellm.store_audit_logs) if is_audit_logging_enabled(store_audit_logs=_store_audit_logs): - print( # noqa: T201 + print_startup_message( f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}" ) else: @@ -7232,9 +7231,7 @@ class ProxyConfig: if model_list: router_params["model_list"] = model_list validate_auto_router_capability_limits(model_list, limit=_license_check.auto_router_capability_limit()) - print( # noqa: T201 - "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" - ) + print_startup_message("\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m") for model in model_list: ### LOAD FROM os.environ/ ### for k, v in model["litellm_params"].items(): @@ -7251,7 +7248,7 @@ class ProxyConfig: complexity_router_config=complexity_router_config, config_file_path=config_file_path, ) - print(f"\033[32m {model.get('model_name', '')}\033[0m") # noqa: T201 + print_startup_message(f"\033[32m {model.get('model_name', '')}\033[0m") litellm_model_name = model["litellm_params"]["model"] litellm_model_api_base = model["litellm_params"].get("api_base", None) if "ollama" in litellm_model_name and litellm_model_api_base is None: @@ -9672,9 +9669,12 @@ async def initialize( master_key, \ user_custom_auth, \ prisma_client - from litellm.proxy.common_utils.banner import show_banner + from litellm.proxy.common_utils.banner import LITELLM_BANNER, show_banner - show_banner() + if litellm.json_logs: + print_startup_message(LITELLM_BANNER) + else: + show_banner() if os.getenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "").lower() != "true": generate_feedback_box() user_model = model @@ -11930,7 +11930,7 @@ class ProxyStartupEvent: and proxy_logging_obj.slack_alerting_instance.alerting is not None and prisma_client is not None ): - print("Alerting: Initializing Weekly/Monthly Spend Reports") # noqa: T201 + print_startup_message("Alerting: Initializing Weekly/Monthly Spend Reports") spend_report_frequency: Final[str] = general_settings.get("spend_report_frequency", "7d") or "7d" days: Final = int(spend_report_frequency[:-1]) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 052c52982b6..262114ee32b 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -987,6 +987,7 @@ def refused_boot_log( overrides: Mapping[str, str], *, config: Path | None = None, + extra_arguments: tuple[str, ...] = (), ) -> str: """Start the proxy and return its log once it exits non-zero instead of becoming ready.""" root: Final = _proxy_root() @@ -1004,6 +1005,7 @@ def refused_boot_log( "--num_workers", "1", *DB_PUSH, + *extra_arguments, ) launch: Final = _launch(command, root, environment, output) try: diff --git a/tests/integration/configuration/test_json_logs_startup.py b/tests/integration/configuration/test_json_logs_startup.py new file mode 100644 index 00000000000..de139b7986f --- /dev/null +++ b/tests/integration/configuration/test_json_logs_startup.py @@ -0,0 +1,397 @@ +"""JSON logging keeps every line the proxy writes during boot, serving and shutdown machine-parseable. + +A deployment that turns on JSON logs ships the proxy's stdout to a log pipeline, and every line that is +not a JSON object breaks that pipeline's parser. This contract boots an owned proxy with JSON logs on, +serves a chat request on every configured model plus one rejected request, stops the proxy and +rejects any line of its output that is not one JSON record, while checking that the +startup information those lines used to carry is still present as JSON records. JSON mode is switched +on through the environment and through the config, at one and two uvicorn workers, at the INFO, the +default and the WARNING log level, under a concurrent burst that lands on both workers, under an operator +--log_config that names the proxy logger or gives it a handler of its own, and with a policy list longer +than the stdout cap. The model list keeps a model name the secret patterns would match. The deprecated +--enforce_prisma_migration_check notice is one JSON record under a config-only json_logs, and a boot whose +config cannot be loaded still prints it before failing. +""" + +from __future__ import annotations + +import re +import uuid +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import pytest +import yaml +from pydantic import JsonValue, TypeAdapter, ValidationError + +from tests.integration._support.client import Gateway +from tests.integration._support.process import owned_proxy_process, refused_boot_log + +JSON_RECORD: Final = TypeAdapter(dict[str, JsonValue]) +PLAIN_LOG_LINE: Final = re.compile(r"^\d\d:\d\d:\d\d - LiteLLM") +PRISMA_ENGINE_TEXT: Final = re.compile(r"query-engine [0-9a-f]{40}") +STARTED_WORKER: Final = "Started server process" +MODEL_LIST_HEADER: Final = "Proxy initialized with Config, Set models:" +CACHE_LINE: Final = "Setting Cache on Proxy" +CONNECTION_BUDGET_LINE: Final = "Postgres connection budget" +REDIS_AUTH_CACHE_LINE: Final = "enable_redis_auth_cache" +BANNER_FRAGMENT: Final = "██╗" +FEEDBACK_BOX_LINE: Final = "Give Feedback / Get Help" +WORKER_COUNTS: Final = (1, 2) +UNKNOWN_MODEL_STATUS: Final = 400 +BURST_SIZE: Final = 16 +CACHE_SETTINGS: Final[Mapping[str, JsonValue]] = {"cache": True, "cache_params": {"type": "redis"}} +STDOUT_CAP: Final = 4096 +POLICY_COUNT: Final = 60 +POLICY_LIST_HEADER: Final = "LiteLLM Policy Engine: Loaded" +TRUNCATION_MARKER: Final = "litellm_truncated" +DEPRECATED_FLAG: Final = "--enforce_prisma_migration_check" +DEPRECATED_FLAG_NOTICE: Final = "--enforce_prisma_migration_check is deprecated and has no effect" +UNPARSEABLE_CONFIG: Final = "model_list: [\n - bad: yaml: here\n" +OPERATOR_LOG_CONFIG: Final[Mapping[str, JsonValue]] = { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"json": {"()": "litellm._logging.JsonFormatter"}}, + "handlers": {"stdout": {"class": "logging.StreamHandler", "stream": "ext://sys.stdout", "formatter": "json"}}, + "loggers": {"LiteLLM Proxy": {"level": "INFO"}}, + "root": {"level": "INFO", "handlers": ["stdout"]}, +} +OPERATOR_SILENCING_LOG_CONFIG: Final[Mapping[str, JsonValue]] = { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"json": {"()": "litellm._logging.JsonFormatter"}}, + "handlers": { + "stdout": {"class": "logging.StreamHandler", "stream": "ext://sys.stdout", "formatter": "json"}, + "silence": {"class": "logging.NullHandler"}, + }, + "loggers": {"LiteLLM Proxy": {"level": "INFO", "handlers": ["silence"], "propagate": False}}, + "root": {"level": "INFO", "handlers": ["stdout"]}, +} + + +def _models() -> tuple[str, str]: + stamp: Final = uuid.uuid4().hex + return (f"json-logs-first-{stamp}", f"json-logs-second-{stamp}") + + +def _policies() -> dict[str, JsonValue]: + stamp: Final = uuid.uuid4().hex + return { + f"startup-listing-policy-{n:02d}-{stamp}": { + "description": f"policy {n:02d} of a listing that runs past the stdout cap meant for request payloads" + } + for n in range(POLICY_COUNT) + } + + +def _write_config( + directory: Path, + upstream_url: str, + models: tuple[str, ...], + litellm_settings: Mapping[str, JsonValue], + extra_sections: Mapping[str, JsonValue] | None = None, +) -> Path: + config: Final = directory / f"json_logs_{uuid.uuid4().hex}.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": model, + "litellm_params": { + "model": f"openai/{model}", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + }, + } + for model in models + ], + "litellm_settings": dict(litellm_settings), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_prisma_schema_update": True, + }, + **(extra_sections or {}), + }, + sort_keys=False, + ) + ) + return config + + +def _record(line: str) -> dict[str, JsonValue] | None: + stripped: Final = line.strip() + if not stripped.startswith("{"): + return None + try: + return JSON_RECORD.validate_json(stripped) + except ValidationError: + return None + + +def _lines(log: str) -> tuple[str, ...]: + return tuple(line for line in PRISMA_ENGINE_TEXT.sub("", log).splitlines() if line.strip()) + + +def _messages(log: str) -> tuple[str, ...]: + records: Final = tuple(record for record in (_record(line) for line in _lines(log)) if record is not None) + return tuple(str(record["message"]) for record in records if "message" in record) + + +def _non_json_lines(log: str) -> tuple[str, ...]: + return tuple(line for line in _lines(log) if _record(line) is None) + + +def _boot_serve_and_stop( + gateway: Gateway, + directory: Path, + config: Path, + overrides: Mapping[str, str], + workers: int, + models: tuple[str, ...], + extra_arguments: tuple[str, ...] = (), +) -> str: + with owned_proxy_process( + gateway, + directory, + overrides, + config=config, + workers=workers, + database_setup=(), + extra_arguments=extra_arguments, + ) as owned: + _serve(owned.gateway, models) + log: Final = owned.log + return log.read_text() + + +def _serve(gateway: Gateway, models: tuple[str, ...]) -> None: + """One chat per model and one rejected request, so the log holds the request-time lines of served + and failed requests next to the boot lines.""" + for model in models: + reply: Final = gateway.chat(model, text=f"json logs {uuid.uuid4().hex}") + assert reply["model"] == model, reply + rejected: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": f"missing-{uuid.uuid4().hex}", "messages": [{"role": "user", "content": "json logs"}]}, + ) + assert rejected.status_code == UNKNOWN_MODEL_STATUS, rejected.text + + +def _assert_every_line_is_json( + log: str, workers: int, models: tuple[str, ...], *, redis_auth_cache_line: bool = True +) -> None: + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + messages: Final = _messages(log) + assert sum(STARTED_WORKER in message for message in messages) == workers, messages + assert any(MODEL_LIST_HEADER in message for message in messages), messages + for model in models: + assert any(model in message for message in messages), (model, messages) + assert any(CACHE_LINE in message for message in messages), messages + if redis_auth_cache_line: + assert any(REDIS_AUTH_CACHE_LINE in message for message in messages), messages + assert any(CONNECTION_BUDGET_LINE in message for message in messages), messages + assert sum(BANNER_FRAGMENT in message for message in messages) == workers, messages + assert sum(FEEDBACK_BOX_LINE in message for message in messages) == workers, messages + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_env_keeps_every_line_json(gateway: Gateway, tmp_path: Path, workers: int) -> None: + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, CACHE_SETTINGS) + log: Final = _boot_serve_and_stop( + gateway, tmp_path, config, {"JSON_LOGS": "true", "LITELLM_LOG": "INFO"}, workers, models + ) + _assert_every_line_is_json(log, workers, models) + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_config_keeps_every_line_json(gateway: Gateway, tmp_path: Path, workers: int) -> None: + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, workers, models) + _assert_every_line_is_json(log, workers, models) + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_with_operator_log_config_keeps_every_line_json( + gateway: Gateway, tmp_path: Path, workers: int +) -> None: + """uvicorn applies an operator `--log_config` after the CLI process switched JSON on and before + `initialize()` switches it on again through the config; a config that names the proxy logger leaves it + without handlers in that window, where the banner and the feedback box are emitted. They still arrive, + as JSON, one per worker, and the operator's own formatter keeps uvicorn's lines JSON. The Redis + auth-cache line is a regular logger line written in that same window; this log_config removed the + proxy logger's handlers and left propagation off, so it is lost at one worker on main and here alike.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log_config: Final = tmp_path / "operator_log_config.yaml" + log_config.write_text(yaml.safe_dump(dict(OPERATOR_LOG_CONFIG))) + log: Final = _boot_serve_and_stop( + gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, workers, models, ("--log_config", str(log_config)) + ) + _assert_every_line_is_json(log, workers, models, redis_auth_cache_line=False) + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_with_operator_handler_on_the_proxy_logger_keeps_the_startup_notices( + gateway: Gateway, tmp_path: Path, workers: int +) -> None: + """An operator `--log_config` that gives the proxy logger a handler of its own, here a NullHandler that + silences it, is in force in the window where the banner and the feedback box are emitted. The prints + they replace wrote to stdout whatever the logging config said, so both still arrive there as JSON, one + per worker. The Redis auth-cache line is a regular logger line written in that same window and goes + to the operator's handler, so it is silenced on main and here alike.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log_config: Final = tmp_path / "operator_silencing_log_config.yaml" + log_config.write_text(yaml.safe_dump(dict(OPERATOR_SILENCING_LOG_CONFIG))) + log: Final = _boot_serve_and_stop( + gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, workers, models, ("--log_config", str(log_config)) + ) + _assert_every_line_is_json(log, workers, models, redis_auth_cache_line=False) + + +def test_json_logs_keeps_a_startup_listing_longer_than_the_stdout_cap_whole(gateway: Gateway, tmp_path: Path) -> None: + """A startup listing longer than MAX_STRING_LENGTH_STDOUT_LOG reaches the pipeline whole. The policy list + stands in for every such listing (the required migration SQL on a drifted database is the other one): + the print it replaces showed it whole, and the stdout cap exists for echoed request payloads, so the JSON + record must not be middle-truncated.""" + models: Final = _models() + policies: Final = _policies() + config: Final = _write_config( + tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}, {"policies": policies} + ) + + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models) + + _assert_every_line_is_json(log, 1, models) + listings: Final = tuple(message for message in _messages(log) if POLICY_LIST_HEADER in message) + assert len(listings) == 1, listings + assert len(listings[0]) > STDOUT_CAP, len(listings[0]) + assert TRUNCATION_MARKER not in listings[0], listings[0] + for policy in policies: + assert policy in listings[0], (policy, listings[0]) + + +def test_json_logs_model_list_keeps_a_model_name_the_secret_patterns_would_match( + gateway: Gateway, tmp_path: Path +) -> None: + """The model list printed every model name as configured. The sk- key pattern matches inside a name such as + `risk-assessment-...`, so the JSON record must carry that name as the print did, not `riREDACTED`.""" + models: Final = (f"risk-assessment-{uuid.uuid4().hex}",) + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models) + + _assert_every_line_is_json(log, 1, models) + assert f" {models[0]}" in _messages(log), _messages(log) + + +def test_json_logs_default_level_keeps_every_line_json(gateway: Gateway, tmp_path: Path) -> None: + """Without LITELLM_LOG the proxy loggers sit at WARNING; the startup notices a plain-mode operator + always sees still arrive, as JSON, and nothing arrives as plaintext.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {}, 1, models) + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + messages: Final = _messages(log) + assert sum(STARTED_WORKER in message for message in messages) == 1, messages + for fragment in ( + BANNER_FRAGMENT, + FEEDBACK_BOX_LINE, + MODEL_LIST_HEADER, + CACHE_LINE, + CONNECTION_BUDGET_LINE, + *models, + ): + assert any(fragment in message for message in messages), (fragment, messages) + + +def test_json_logs_warning_level_keeps_startup_notices(gateway: Gateway, tmp_path: Path) -> None: + """`LITELLM_LOG=WARNING` drops the INFO logger lines; the startup notices a plain-mode operator sees at + every level still arrive, as JSON, and nothing arrives as plaintext.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "WARNING"}, 1, models) + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + messages: Final = _messages(log) + assert not any(REDIS_AUTH_CACHE_LINE in message for message in messages), messages + for fragment in ( + BANNER_FRAGMENT, + FEEDBACK_BOX_LINE, + MODEL_LIST_HEADER, + CACHE_LINE, + CONNECTION_BUDGET_LINE, + *models, + ): + assert any(fragment in message for message in messages), (fragment, messages) + + +def test_json_logs_burst_on_two_workers_keeps_every_line_json(gateway: Gateway, tmp_path: Path) -> None: + """Concurrent requests spread over both workers interleave their log writes; every line must still be + one complete JSON record.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + with owned_proxy_process( + gateway, tmp_path, {"LITELLM_LOG": "INFO"}, config=config, workers=2, database_setup=() + ) as owned: + with ThreadPoolExecutor(max_workers=BURST_SIZE) as pool: + replies: Final = tuple( + pool.map(lambda index: owned.gateway.chat(models[index % 2], text=f"burst {index}"), range(BURST_SIZE)) + ) + assert tuple(reply["model"] for reply in replies) == tuple(models[index % 2] for index in range(BURST_SIZE)) + log: Final = owned.log + text: Final = log.read_text() + plaintext: Final = _non_json_lines(text) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + assert sum(STARTED_WORKER in message for message in _messages(text)) == 2, _messages(text) + + +def test_json_logs_config_makes_the_deprecated_flag_notice_one_json_record(gateway: Gateway, tmp_path: Path) -> None: + """The notice waits until the config is read, so a json_logs set only there makes it a JSON record.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {"json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models, (DEPRECATED_FLAG,)) + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + assert sum(DEPRECATED_FLAG_NOTICE in message for message in _messages(log)) == 1, _messages(log) + + +@pytest.mark.parametrize("config_text", (None, UNPARSEABLE_CONFIG), ids=("missing", "unparseable")) +def test_a_config_that_fails_to_load_still_prints_the_deprecated_flag_notice( + gateway: Gateway, tmp_path: Path, config_text: str | None +) -> None: + """A boot with the deprecated flag and a config file that does not exist or does not parse exits + non-zero, and it still prints the notice once before the error, as it did before the notice waited + for the config.""" + config: Final = tmp_path / f"unloadable_{uuid.uuid4().hex}.yaml" + if config_text is not None: + config.write_text(config_text) + log: Final = refused_boot_log(gateway, tmp_path, {}, config=config, extra_arguments=(DEPRECATED_FLAG,)) + assert log.count(DEPRECATED_FLAG_NOTICE) == 1, log + + +def test_plain_logs_control_prints_startup_text(gateway: Gateway, tmp_path: Path) -> None: + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, CACHE_SETTINGS) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models) + plaintext: Final = _non_json_lines(log) + assert any(BANNER_FRAGMENT in line for line in plaintext), log + assert any(FEEDBACK_BOX_LINE in line for line in plaintext), log + assert any(CACHE_LINE in line for line in plaintext), log + assert any(MODEL_LIST_HEADER in line for line in plaintext), log + for model in models: + assert any(line.strip().endswith(model) or line.strip().endswith(f"{model}\x1b[0m") for line in plaintext), ( + model, + log, + ) + assert any(PLAIN_LOG_LINE.match(line) and REDIS_AUTH_CACHE_LINE in line for line in plaintext), log + assert any(line.startswith(f"LiteLLM Proxy: {CONNECTION_BUDGET_LINE}") for line in plaintext), log diff --git a/tests/unit/proxy/db/test_check_migration.py b/tests/unit/proxy/db/test_check_migration.py index 74a6841cb23..b5685676941 100644 --- a/tests/unit/proxy/db/test_check_migration.py +++ b/tests/unit/proxy/db/test_check_migration.py @@ -1,6 +1,12 @@ +import json +import logging +import sys +from typing import Final import pytest +import litellm +from litellm._logging import turn_on_json, verbose_logger, verbose_proxy_logger def test_check_migration_out_of_sync(mocker): @@ -69,3 +75,33 @@ def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(m monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + + +def test_a_skipped_migration_diff_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """ + With json_logs on, the notice that the migration diff was skipped is one INFO record on stdout, + also while the loggers sit at WARNING, as they do at boot without LITELLM_LOG. + """ + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_logger.setLevel(logging.WARNING) + verbose_proxy_logger.setLevel(logging.WARNING) + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert (record["level"], record["message"]) == ( + "INFO", + "Skipping the migration diff: litellm-proxy-extras has no Prisma runner. " + "Error: import of litellm_proxy_extras.prisma_toolchain halted; None in sys.modules", + ) diff --git a/tests/unit/proxy/policy_engine/test_init_policies.py b/tests/unit/proxy/policy_engine/test_init_policies.py new file mode 100644 index 00000000000..b9b5f278252 --- /dev/null +++ b/tests/unit/proxy/policy_engine/test_init_policies.py @@ -0,0 +1,116 @@ +"""The policy table printed at boot keeps its terminal form in plain mode and is JSON when json_logs is on.""" + +import json +import sys +from collections.abc import Mapping, Sequence +from typing import Final + +import pytest + +import litellm +from litellm._logging import turn_on_json +from litellm.proxy.policy_engine.init_policies import _print_policies_on_startup + +POLICIES: Final[Mapping[str, Mapping[str, object]]] = { + "pii-policy": { + "inherit": "base", + "description": "Block PII", + "guardrails": {"add": ["pii-guard"], "remove": ["noop"]}, + "condition": {"model": "gpt-4o"}, + } +} +ATTACHMENTS: Final[Sequence[Mapping[str, object]]] = ({"policy": "pii-policy", "scope": "*", "teams": ["team-a"]},) +POLICY_ROWS: Final = ( + " - pii-policy (inherits: base)", + " description: Block PII", + " guardrails.add: ['pii-guard']", + " guardrails.remove: ['noop']", + " condition.model: gpt-4o", +) +PLAIN_POLICY_TABLE: Final = ( + "\x1b[92m\nLiteLLM Policy Engine: Loaded 1 policies\x1b[0m\n\n" + "\x1b[94m - pii-policy (inherits: base)\x1b[0m\n" + " description: Block PII\n" + " guardrails.add: ['pii-guard']\n" + " guardrails.remove: ['noop']\n" + " condition.model: gpt-4o\n" +) +NO_ATTACHMENTS_WARNING: Final = ( + "Warning: No policy_attachments configured. Policies will not be applied to any requests." +) + + +@pytest.mark.parametrize( + ("attachments", "expected"), + [ + ( + ATTACHMENTS, + PLAIN_POLICY_TABLE + + "\n\x1b[93mPolicy Attachments: 1 attachment(s)\x1b[0m\n" + + " - pii-policy -> scope=* (global), teams=['team-a']\n\n", + ), + (None, PLAIN_POLICY_TABLE + f"\n\x1b[93m{NO_ATTACHMENTS_WARNING}\x1b[0m\n\n"), + ], + ids=["with-attachments", "without-attachments"], +) +def test_policy_table_prints_unchanged_in_plain_mode( + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], + attachments: Sequence[Mapping[str, object]] | None, + expected: str, +) -> None: + """Operators without JSON logs keep the colored policy table byte for byte.""" + monkeypatch.setattr(litellm, "json_logs", False) + + _print_policies_on_startup(POLICIES, attachments) + + assert capfd.readouterr() == (expected, "") + + +def test_policy_table_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """With JSON logs on the whole table, attachments included, is one INFO record without color codes.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + _print_policies_on_startup(POLICIES, ATTACHMENTS) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["level"] == "INFO" + assert record["message"] == "\n".join( + ( + "LiteLLM Policy Engine: Loaded 1 policies", + "", + *POLICY_ROWS, + "", + "Policy Attachments: 1 attachment(s)", + " - pii-policy -> scope=* (global), teams=['team-a']", + ) + ) + + +def test_missing_attachments_warning_is_a_second_json_record( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """Without attachments the table and the yellow warning are two INFO records on stdout, where the two + prints put them.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + _print_policies_on_startup(POLICIES, None) + + out, err = capfd.readouterr() + assert err == "" + records: Final = [json.loads(raw) for raw in out.splitlines() if raw.strip()] + assert [record["message"] for record in records] == [ + "\n".join(("LiteLLM Policy Engine: Loaded 1 policies", "", *POLICY_ROWS)), + NO_ATTACHMENTS_WARNING, + ] + assert {record["level"] for record in records} == {"INFO"} diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 89a35923be0..d77731aa83b 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -1,4 +1,5 @@ import inspect +import json import logging import os from contextlib import nullcontext @@ -9,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import click import fastapi +import httpx import pytest @@ -23,6 +25,8 @@ import yaml from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string +import litellm +from litellm._logging import turn_on_json from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS from litellm.proxy import proxy_cli from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server @@ -2875,6 +2879,26 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert ("--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out) is warned + def test_the_retired_enforce_prisma_migration_check_opt_in_warns_on_health( + self, capsys: pytest.CaptureFixture[str] + ) -> None: + """`--health` returns before the config is read, where a server start emits the deprecation line; + it still prints it, as it did when the notice came before that return.""" + with patch("httpx.get", return_value=httpx.Response(200, json={})) as mock_get: + run_server.main(["--health", "--enforce_prisma_migration_check"], standalone_mode=False) + + mock_get.assert_called_once_with(url="http://0.0.0.0:4000/health") + assert "--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out + + def test_the_retired_enforce_prisma_migration_check_opt_in_warns_on_test( + self, capsys: pytest.CaptureFixture[str] + ) -> None: + """`--test` returns before the config is read too, and prints the deprecation line before its request.""" + with pytest.raises(ValueError, match="Invalid test value"): + run_server.main(["--test", "--enforce_prisma_migration_check"], standalone_mode=False) + + assert "--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out + @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True) @@ -3935,3 +3959,117 @@ class TestValidateConfigFlag: assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "config OK" in result.output mock_popen.assert_not_called() + + +def _boot_from_config_file(config_path: Path, *flags: str) -> str | None: + env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "JSON_LOGS")} + with patch.dict(os.environ, env, clear=True): + run_server.main(["--config", str(config_path), "--skip_server_startup", *flags], standalone_mode=False) + return os.environ.get("JSON_LOGS") + + +class TestJsonLogsConfigReachesSpawnedWorkers: + """`litellm_settings.json_logs: true` turns JSON on inside the CLI process, but the workers that + uvicorn or granian spawn are fresh interpreters that pick plain or JSON from the JSON_LOGS + environment variable at import. The CLI has to export it, or every worker logs plain text until + its own load_config reaches the key.""" + + def test_json_logs_true_exports_json_logs_to_spawned_workers( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {"json_logs": True}})) + + assert _boot_from_config_file(config_path) == "true" + + def test_plain_logs_config_leaves_json_logs_unset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {}})) + + assert _boot_from_config_file(config_path) is None + + +class TestEnforcePrismaMigrationCheckNoticeAroundTheConfigRead: + """With a config file, the `--enforce_prisma_migration_check` deprecation notice waits until the config + is read, so a `json_logs: true` set only there makes it a JSON record. It is printed once either way, + and a boot whose config read fails still prints it before the error, as it did before the wait.""" + + NOTICE: Final = "--enforce_prisma_migration_check is deprecated and has no effect" + FLAG: Final = "--enforce_prisma_migration_check" + + def test_a_config_read_that_fails_still_prints_the_notice( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + + with pytest.raises(Exception, match="Config file not found"): + _boot_from_config_file(tmp_path / "missing.yaml", self.FLAG) + + assert capsys.readouterr().out.count(self.NOTICE) == 1 + + def test_a_plain_config_prints_the_notice_once( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {}})) + + _boot_from_config_file(config_path, self.FLAG) + + assert capsys.readouterr().out.count(self.NOTICE) == 1 + + def test_a_config_only_json_logs_makes_the_notice_one_json_record( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {"json_logs": True}})) + + _boot_from_config_file(config_path, self.FLAG) + + notice_lines: Final = [line for line in capsys.readouterr().out.splitlines() if self.NOTICE in line] + assert len(notice_lines) == 1, notice_lines + assert self.NOTICE in json.loads(notice_lines[0])["message"] + + +def test_granian_flag_notices_are_json_records_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The yellow notices for the flags Granian ignores (`--max_requests_before_restart`, `--ciphers`) are + printed while the launcher starts; with json_logs on they are INFO records on stdout like every other + line, not plaintext between the records.""" + pytest.importorskip("granian") + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + with ( + patch("granian.Granian") as mock_granian_cls, + patch("granian.constants.Interfaces", SimpleNamespace(ASGI="asgi")), + ): + ProxyInitializationHelpers._init_granian_server( + host="0.0.0.0", + port=4000, + num_workers=2, + ssl_certfile_path=None, + ssl_keyfile_path=None, + max_requests_before_restart=100, + ciphers="ECDHE", + granian_runtime_threads=None, + ) + + mock_granian_cls.return_value.serve.assert_called_once() + out, err = capfd.readouterr() + assert err == "" + records: Final = [json.loads(raw) for raw in out.splitlines() if raw.strip()] + assert [record["message"] for record in records] == [ + "LiteLLM Proxy: Starting server on 0.0.0.0:4000 using Granian", + "LiteLLM: --max_requests_before_restart is not supported by Granian " + "(Granian uses workers_lifetime in seconds, not a per-request limit).", + "LiteLLM: --ciphers is not applied when using --run_granian.", + ] + assert {record["level"] for record in records} == {"INFO"} diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 8c62b225992..55d1587669a 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -7,6 +7,7 @@ import os import re import socket import subprocess +import sys import time import types from datetime import datetime, timedelta, timezone @@ -28,6 +29,7 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as proxy_server_module from litellm._internal_context import current_service_target +from litellm._logging import turn_on_json, verbose_proxy_logger from litellm.caching.caching import RedisCache from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cluster_cache import RedisClusterCache @@ -46,6 +48,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.proxy.common_utils.banner import LITELLM_BANNER from tests._master_key import MASTER_KEY PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" @@ -14266,7 +14269,6 @@ async def test_moderations_response_carries_litellm_call_id_header(): async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplog): """LIT-7836: the /v1/moderations error line must carry the litellm_call_id the client sent, rendered in the message and as a structured log record field.""" - from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException call_id = "moderations-call-7836" @@ -15586,3 +15588,61 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is call(api_provider="openai", capture_rate=0.97), call(api_provider="openai", capture_rate=None), ] + + +def test_feedback_box_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The feedback box is a dozen terminal lines; with JSON logs on it has to land as one parseable + record, not a plaintext line per row.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + proxy_server_module.generate_feedback_box() + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert "Give Feedback / Get Help: https://github.com/BerriAI/litellm/issues/new" in record["message"] + assert "Thank you for using LiteLLM! - Krrish & Ishaan" in record["message"] + assert "\x1b[" not in record["message"] + + +@pytest.mark.asyncio +async def test_initialize_prints_the_banner_between_blank_lines_in_plain_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """Operators without JSON logs see the banner exactly as before, through show_banner()'s flushing + click.echo.""" + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "true") + + await proxy_server_module.initialize() + + out, err = capfd.readouterr() + assert err == "" + assert out.startswith(f"\n{LITELLM_BANNER}\n\n"), out + + +@pytest.mark.asyncio +async def test_initialize_prints_the_banner_as_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """With JSON logs on, the six banner rows are one INFO record with their indentation kept, so the + first line a worker writes parses like every other.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + monkeypatch.setenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "true") + turn_on_json() + + await proxy_server_module.initialize() + + out, err = capfd.readouterr() + assert err == "" + records: Final = [json.loads(raw) for raw in out.splitlines() if raw.strip()] + banner_records: Final = [record for record in records if record["message"] == LITELLM_BANNER] + assert len(banner_records) == 1, records + assert banner_records[0]["level"] == "INFO" diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index cb4d3f82611..36be864a6f7 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -34,6 +34,7 @@ from litellm._logging import ( _parse_json_logs_env, _plain_log_format, _stdout_truncation_marker, + print_startup_message, turn_on_json, format_base64_size, session_id_var, @@ -1824,3 +1825,222 @@ def test_print_verbose_formats_args_only_when_set_verbose_is_on(monkeypatch, cap logging_module.print_verbose("key %s value %s", "k1", _StrCountingValue()) assert _StrCountingValue.str_calls == 1 assert capsys.readouterr().out == "key k1 value cached-vector\n" + + +def test_print_startup_message_prints_plain_mode_text_unchanged( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """Operators without JSON logs keep the colored, spaced terminal output exactly as before.""" + monkeypatch.setattr(litellm, "json_logs", False) + + print_startup_message("\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n") + + out, err = capfd.readouterr() + assert out == "\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n\n" + assert err == "" + + +def test_print_startup_message_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """With JSON logs on, a multi-line colored startup message is one INFO record with the colors and the + surrounding blank lines removed, so a log pipeline parses it like every other line. The logger sits + at WARNING, as it does in the CLI process, and the record still goes out: a print has no level.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.WARNING) + + print_startup_message( + "\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n\033[32m gpt-4o\033[0m\n\n" + ) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "LiteLLM: Proxy initialized with Config, Set models:\n gpt-4o" + assert record["level"] == "INFO" + assert record["component"] == "LiteLLM Proxy" + assert str(record["logger"]).startswith("test_logging.py:"), record["logger"] + + +def test_print_startup_message_yellow_notice_is_an_info_record_on_stdout( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """An operator notice printed in yellow (an ignored flag, a deprecation) stays on stdout in JSON mode, + where the print put it, as one INFO record with the color codes removed.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.WARNING) + + print_startup_message("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "LiteLLM: --ciphers is not applied when using --run_granian." + assert record["level"] == "INFO" + + +def test_print_startup_message_shows_above_the_json_handler_level( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """`LITELLM_LOG=WARNING` sets the JSON handler to WARNING, which drops the INFO logger lines; the + startup notices still go out, as the prints they replace did at every level.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.INFO) + for json_handler in verbose_proxy_logger.handlers: + monkeypatch.setattr(json_handler, "level", logging.WARNING) + + print_startup_message("\033[94m\nSetting Cache on Proxy\033[0m") + verbose_proxy_logger.info("enable_redis_auth_cache is not set") + + out, _ = capfd.readouterr() + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert [json.loads(raw)["message"] for raw in lines] == ["Setting Cache on Proxy"] + + +def test_print_startup_message_emits_json_after_an_operator_log_config_removed_the_handlers( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """uvicorn applies an operator `--log_config` after proxy_cli's turn_on_json() and before initialize() + re-installs the JSON handler through load_config(); a config that names "LiteLLM Proxy" leaves the + logger without handlers in that window, where the banner and the feedback box are emitted, and they + still go out as one JSON record each, on stdout.""" + import logging.config + + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + request.addfinalizer(turn_on_json) + logging.config.dictConfig( + {"version": 1, "disable_existing_loggers": False, "loggers": {"LiteLLM Proxy": {"level": "INFO"}}} + ) + assert verbose_proxy_logger.handlers == [] + + print_startup_message("\n\033[1;37m# Give Feedback / Get Help #\033[0m\n") + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "# Give Feedback / Get Help #" + assert record["level"] == "INFO" + assert record["component"] == "LiteLLM Proxy" + + +OPERATOR_HANDLERS_ON_THE_PROXY_LOGGER: Final = { + "silenced": {"class": "logging.NullHandler"}, + "stderr_warning": { + "class": "logging.StreamHandler", + "stream": "ext://sys.stderr", + "level": "WARNING", + "formatter": "plain", + }, +} + + +@pytest.mark.parametrize("handler", sorted(OPERATOR_HANDLERS_ON_THE_PROXY_LOGGER)) +def test_print_startup_message_stays_on_stdout_under_an_operator_handler_on_the_proxy_logger( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str], handler: str +) -> None: + """An operator `--log_config` that gives "LiteLLM Proxy" a handler of its own, a NullHandler that silences + it or a plain WARNING-level handler on stderr, is in force when the banner and the feedback box are + emitted; the print they replace wrote to stdout whatever the logging config said, so the record still + goes out as one JSON record on stdout and the operator's handler gets nothing.""" + import logging.config + + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + request.addfinalizer(turn_on_json) + logging.config.dictConfig( + { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"plain": {"format": "%(levelname)s %(message)s"}}, + "handlers": {"proxy": OPERATOR_HANDLERS_ON_THE_PROXY_LOGGER[handler]}, + "loggers": {"LiteLLM Proxy": {"level": "INFO", "handlers": ["proxy"], "propagate": False}}, + } + ) + assert len(verbose_proxy_logger.handlers) == 1 and not isinstance( + verbose_proxy_logger.handlers[0], LevelRoutingStreamHandler + ) + + print_startup_message("\n\033[1;37m# Give Feedback / Get Help #\033[0m\n") + + out, err = capfd.readouterr() + assert err == "", err + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "# Give Feedback / Get Help #" + assert record["level"] == "INFO" + assert record["component"] == "LiteLLM Proxy" + + +def test_print_startup_message_keeps_a_listing_longer_than_the_stdout_cap_whole( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The required migration SQL and the policy list run past MAX_STRING_LENGTH_STDOUT_LOG on a drifted + deployment. The prints they replace showed them whole, and that cap exists for echoed request payloads, + so the JSON record carries the whole listing too.""" + monkeypatch.setenv("MAX_STRING_LENGTH_STDOUT_LOG", "500") + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + commands: Final = tuple(f'ALTER TABLE "LiteLLM_Table{n}" ADD COLUMN "column_{n}" TEXT;' for n in range(40)) + listing: Final = "\n".join(("Changes to DB Schema detected", "Required SQL commands:", *commands)) + assert len(listing) > 500 + + print_startup_message(listing) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert LITELLM_TRUNCATED_PAYLOAD_FIELD not in record["message"] + assert record["message"] == listing + + +def test_print_startup_message_keeps_a_model_name_the_secret_patterns_would_match( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The model list printed `risk-assessment-gpt4o` as configured; the sk- key pattern matches inside + that name, so the JSON record has to carry the name as the print did, not `riREDACTED`.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + assert secret_redaction.redact_string("risk-assessment-gpt4o") != "risk-assessment-gpt4o" + + print_startup_message("\033[32m risk-assessment-gpt4o\033[0m") + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + assert json.loads(lines[0])["message"] == " risk-assessment-gpt4o" + + +def test_print_startup_message_blank_spacer_emits_no_json_record( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """A spacer that only pads the terminal has nothing to say to a log pipeline.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.INFO) + + print_startup_message("") + print_startup_message("\n\033[0m\n") + + assert capfd.readouterr() == ("", "") From 8f31c7171dfbd28c53d294121a28c5d232d9137c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 15:12:53 -0700 Subject: [PATCH 025/147] fix(router): keep a lone deployment warm under router allowed_fails unless the policy names the exception (#48644) * fix(router): keep a lone deployment warm under router allowed_fails unless the policy names the exception Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): give cooldown fixtures a sibling deployment so allowed_fails=0 still forces a cooldown * fix(router): keep benching a lone deployment that has a fallback group or an empty deployment policy * fix(router): accept dictionary fallback targets in the lone deployment check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_utils/cooldown_handlers.py | 44 ++++++- .../routing/test_payment_required_cooldown.py | 74 +++++++++++ .../router_utils/test_cooldown_handlers.py | 119 ++++++++++++++---- .../test_fallback_event_handlers.py | 14 ++- tests/unit/test_router/test_router.py | 14 ++- 5 files changed, 234 insertions(+), 31 deletions(-) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index c362c6899db..330500789a4 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -185,7 +185,14 @@ def _should_cooldown_based_on_deployment_policy( override that safety net, matching `_has_explicit_allowed_fails_policy_for_exception`. """ allowed_fails_from_policy: Final = _resolve_allowed_fails_from_policy(dep_policy, original_exception) - if allowed_fails_from_policy is None and dep_allowed_fails is not None and is_single_deployment_model_group: + if ( + allowed_fails_from_policy is None + and is_single_deployment_model_group + and ( + dep_allowed_fails is not None + or litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None + ) + ): return False allowed_fails_override: Final[int | None] = ( @@ -399,6 +406,34 @@ def _should_run_cooldown_logic( return True +_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...]) + + +def _fallback_target_name(target: str | Mapping[str, object]) -> str: + return target if isinstance(target, str) else str(target.get("model", "")) + + +def _fallback_chain_names( + entry: Mapping[str, Sequence[str | Mapping[str, object]]], +) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()} + + +def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool: + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group + + configured: Final[object] = getattr(litellm_router_instance, "fallbacks", None) + if requested_model_group is None or not configured: + return False + chains: Final = tuple( + _fallback_chain_names(entry) + for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) + if isinstance(entry, Mapping) + ) + fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group) + return bool(fallback_group) + + def _should_cooldown_deployment( litellm_router_instance: LitellmRouter, deployment: str, @@ -431,6 +466,9 @@ def _should_cooldown_deployment( litellm_router_instance.routing_group_has_alternatives(requested_model_group) or litellm_router_instance.team_model_has_alternatives(deployment) ) + is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group( + litellm_router_instance, requested_model_group + ) exception_status_int: Final = cast_exception_status_to_int(exception_status) dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment) @@ -450,7 +488,7 @@ def _should_cooldown_deployment( original_exception, dep_policy, dep_allowed_fails, - is_single_deployment_model_group, + is_lone_deployment, ) ## BASE CASE - single deployment @@ -495,6 +533,8 @@ def _should_cooldown_deployment( return False else: + if is_lone_deployment and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None: + return False return should_cooldown_based_on_allowed_fails_policy( litellm_router_instance=litellm_router_instance, deployment=deployment, diff --git a/tests/integration/routing/test_payment_required_cooldown.py b/tests/integration/routing/test_payment_required_cooldown.py index bc5cf1a5913..865edf31260 100644 --- a/tests/integration/routing/test_payment_required_cooldown.py +++ b/tests/integration/routing/test_payment_required_cooldown.py @@ -1,5 +1,6 @@ from __future__ import annotations +import itertools import json import re import uuid @@ -390,6 +391,79 @@ def test_a_router_allowed_fails_policy_cools_a_lone_402_deployment_past_its_coun assert markers_received(refusing) == (first.marker, second.marker) +def failing_reply(request: Request) -> Reply: + return model_list_reply() if is_model_info_probe(request) else anthropic_error(500, "upstream down") + + +def assert_failed(outcome: Outcome) -> None: + assert outcome.status == 500, (outcome.status, outcome.detail) + assert "upstream down" in outcome.detail, outcome.detail + + +@pytest.mark.timeout(900) +def test_a_router_allowed_fails_count_leaves_a_lone_deployment_warm_but_benches_one_with_a_sibling_or_fallback( + tmp_path: Path, +) -> None: + lone_group: Final = new_group() + pair_group: Final = new_group() + primary_group: Final = new_group() + backup_group: Final = new_group() + with ExitStack() as stack: + upstream: Final = stack.enter_context(gateway_from_environment()) + cache: Final = stack.enter_context(owned_redis(tmp_path)) + settings: Final = router_settings(cache, allowed_fails=1, fallbacks=[{primary_group: [backup_group]}]) + owned: Final = stack.enter_context( + owned_proxy_process( + upstream, + tmp_path, + {}, + config=cooldown_config(tmp_path, settings), + remove_environment=("PROMETHEUS_MULTIPROC_DIR",), + ) + ) + failing: Final = stack.enter_context(wire_server(failing_reply)) + healthy: Final = stack.enter_context(wire_server(healthy_reply)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + lone_id: Final = deployment(scenario, lone_group, anthropic_params(failing)) + failed: Final = tuple(raw_chat(owned.gateway, lone_group, new_marker()) for _ in range(4)) + for outcome in failed: + assert_failed(outcome) + assert not cooled(cache, lone_id) + assert markers_received(failing) == markers(failed) + + paired_id: Final = two_deployments(scenario, pair_group, failing, healthy) + history: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, paired_id), + (raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, paired_id), [(outcome.status, outcome.detail[:80]) for outcome in history] + assert sum(outcome.status == 500 for outcome in history) == 1 + served: Final = tuple(raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(10)) + assert_served(served, chat_identity) + assert markers_received(healthy) == (*markers(o for o in history if o.status == 200), *markers(served)) + + primary_id: Final = deployment(scenario, primary_group, anthropic_params(failing)) + deployment(scenario, backup_group, openai_params(healthy)) + failing.drain() + healthy.drain() + fell_back: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, primary_id), + (raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, primary_id) and fell_back, [(o.status, o.detail[:80]) for o in fell_back] + assert_served(fell_back, chat_identity) + backed: Final = tuple(raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(4)) + assert_served(backed, chat_identity) + received: Final = markers_received(failing) + assert received[: len(fell_back)] == markers(fell_back) + assert len(received) <= len(fell_back) + 1, received + assert {*markers(fell_back), *markers(backed)} <= set(markers_received(healthy)) + + def test_a_deployment_allowed_fails_policy_cools_a_lone_402_deployment(rig: CooldownRig) -> None: group: Final = new_group() policy: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}} diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 0c9d7ab9f85..bab6d257579 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -2931,10 +2931,65 @@ async def test_bad_request_error_does_not_cool_down_a_deployment( assert tuple(cooldown_deployments) == () -@pytest.mark.parametrize("use_policy", [False, True], ids=["allowed_fails", "allowed_fails_policy"]) +TIMED_OUT_DEPLOYMENT: Final = { + "model_name": "test-model", + "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, + "model_info": {"id": "test-deployment"}, +} +BACKUP_DEPLOYMENT: Final = { + "model_name": "backup-model", + "litellm_params": {"model": "openai/backup-model", "api_key": "test-key", "api_base": "https://backup.test/v1"}, + "model_info": {"id": "backup-deployment"}, +} +BACKUP_REPLY: Final = { + "id": "chatcmpl-backup", + "object": "chat.completion", + "created": 1700000000, + "model": "backup-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +async def _lone_call_outcome(router: Router) -> type[Exception]: + try: + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "timeout test"}]) + except (litellm.Timeout, RouterRateLimitError) as error: + return type(error) + pytest.fail("the timed-out deployment answered") + + +BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, RouterRateLimitError, RouterRateLimitError) + + +@pytest.mark.parametrize( + ("router_kwargs", "model_info", "outcomes"), + [ + ({"allowed_fails": 1}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails": 1}, {"allowed_fails_policy": {}}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=1)}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)}, {}, BENCHED), + ({"allowed_fails": 1, "fallbacks": [{"test-model": ["backup-model"]}]}, {}, BENCHED), + ( + {"allowed_fails": 1, "fallbacks": [{"test-model": [{"model": "backup-model", "temperature": 0}]}]}, + {}, + BENCHED, + ), + ], + ids=[ + "allowed_fails", + "allowed_fails_with_an_empty_deployment_policy", + "policy_names_another_exception", + "policy_names_timeout", + "allowed_fails_with_a_fallback_group", + "allowed_fails_with_a_dict_fallback_target", + ], +) @pytest.mark.asyncio -async def test_allowed_fails_cools_a_single_deployment( - use_policy: bool, +async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_policy_names_the_exception( + router_kwargs: Mapping[str, object], + model_info: Mapping[str, object], + outcomes: tuple[type[Exception], ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2942,38 +2997,52 @@ async def test_allowed_fails_cools_a_single_deployment( route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=httpx.ReadTimeout("upstream timeout") ) - policy_kwargs: Final = ( - {"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)} - if use_policy - else {"allowed_fails": 1} + deployment: Final = {**TIMED_OUT_DEPLOYMENT, "model_info": {**TIMED_OUT_DEPLOYMENT["model_info"], **model_info}} + router: Final = Router(model_list=[deployment], cooldown_time=300, num_retries=0, **router_kwargs) + + seen: Final = tuple([await _lone_call_outcome(router) for _ in range(5)]) + + assert seen == outcomes + assert len(route.calls) == outcomes.count(litellm.Timeout) + + +async def _group_call_outcome(router: Router) -> str: + try: + response: Final = await router.acompletion( + model="test-model", messages=[{"role": "user", "content": "timeout test"}] + ) + except litellm.Timeout: + return "timeout" + return response.id + + +@pytest.mark.asyncio +async def test_router_allowed_fails_still_benches_a_deployment_that_has_a_sibling_in_its_group( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + timed_out: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=httpx.ReadTimeout("upstream timeout") ) + respx_mock.post("https://backup.test/v1/chat/completions").mock(return_value=httpx.Response(200, json=BACKUP_REPLY)) router: Final = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, - "model_info": {"id": "test-deployment"}, - } - ], + model_list=[TIMED_OUT_DEPLOYMENT, {**BACKUP_DEPLOYMENT, "model_name": "test-model"}], cooldown_time=300, num_retries=0, - **policy_kwargs, + allowed_fails=1, ) - for _ in range(2): - with pytest.raises(litellm.Timeout): - await router.acompletion( - model="test-model", - messages=[{"role": "user", "content": "timeout test"}], - ) - + while len(timed_out.calls) < 2: + await _group_call_outcome(router) cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, - parent_otel_span=None, + litellm_router_instance=router, parent_otel_span=None ) + served: Final = tuple([await _group_call_outcome(router) for _ in range(5)]) - assert len(route.calls) == 2 assert tuple(cooldown_deployments) == ("test-deployment",) + assert served == ("chatcmpl-backup",) * 5 + assert len(timed_out.calls) == 2 def test_per_user_session_upstream_errors_never_cool_down_the_shared_deployment(single_deployment_router): diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 53375ee0ab4..c5a476c450b 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -1255,7 +1255,12 @@ class TestTriggerCooldownForFailedDeployment: "model_name": "fallback-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "fallback-deployment"}, - } + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-sibling"}, + }, ], allowed_fails=0, cooldown_time=60, @@ -1345,7 +1350,12 @@ async def test_a_fallback_deployment_whose_stream_stalled_stays_out_of_cooldown( "model_name": "fallback-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "fallback-deployment"}, - } + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-sibling"}, + }, ], allowed_fails=0, cooldown_time=60, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e2ee80b7874..d9b655a0aff 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -12212,7 +12212,12 @@ class TestBackgroundResponseCostPollCooldown: "model_name": "gpt-4.1", "litellm_params": {"model": "openai/gpt-4.1"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, ) @@ -12340,7 +12345,12 @@ class TestCallerTimeoutCooldown: "model_name": "slow-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "slow-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, cooldown_time=120, From cb70bc560e1bbd8a1d619fb5dead7f2f0c333e74 Mon Sep 17 00:00:00 2001 From: Ambuj Upadhyay Date: Wed, 28 Oct 2026 03:51:11 +0530 Subject: [PATCH 026/147] fix(streaming): dict-usage arm drops prompt_tokens_details/completion_tokens_details its sibling arms preserve (#36678) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(streaming): dict-usage arm no longer drops token detail fields In _dispatch_provider_chunk's openai/azure else-branch, a plain-dict usage was rebuilt into litellm.Usage from only prompt_tokens / completion_tokens / total_tokens. prompt_tokens_details (cached_tokens), completion_tokens_details (reasoning_tokens) and cache token counts were silently dropped — while both sibling arms (Usage passthrough, BaseModel via Usage(**model_dump())) preserve them. Every in-repo chunk producer currently normalizes usage to Usage/BaseModel before this point, so the dict arm is reachable today only via chunks litellm did not construct (e.g. a custom streaming_decoder feeding CustomStreamWrapper). When it does fire on a usage-only final chunk (stream_options.include_usage), the stripped Usage is exactly what the stream consumer receives — the early-return path has no later aggregation to repair it. Fix: pass the full dict through Usage(**usage_dict). Usage.__init__ coerces nested detail dicts into their wrapper types and tolerates extra keys. Side effect: a legitimate 0 count is no longer coerced to None by the old 'or None' expressions. Co-Authored-By: Claude Fable 5 * chore: drop explanatory source comment per repo comment policy Rationale lives in the PR body and the regression tests. Co-Authored-By: Claude Fable 5 * fix(streaming): type the dict-usage unpack for the basedpyright gate Usage(**response_obj["usage"]) unpacks a dict whose value type pyright cannot see, and the delta gate reports that once per typed constructor parameter: five, against the three .get() calls the line replaced. Moving the unpack behind _usage_from_payload keeps the same Usage(**payload) call and is reported once, at the call site, so the rule count ends below its merge base --------- Co-authored-by: Claude Fable 5 --- .../litellm_core_utils/streaming_handler.py | 11 ++- .../test_streaming_handler.py | 99 +++++++++++++++++++ 2 files changed, 105 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 8fb602ebeb0..b64e146ca08 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1589,11 +1589,7 @@ class CustomStreamWrapper: setattr( model_response, "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None, - completion_tokens=response_obj["usage"].get("completion_tokens", None) or None, - total_tokens=response_obj["usage"].get("total_tokens", None) or None, - ), + _usage_from_payload(response_obj["usage"]), ) elif isinstance(response_obj["usage"], Usage): setattr( @@ -2662,6 +2658,11 @@ class CustomStreamWrapper: _strip_sse_data_from_chunk = strip_sse_data_from_chunk +def _usage_from_payload(payload: Mapping[str, Any]) -> Usage: + """Build ``Usage`` from a wire-shaped usage dict, forwarding every key.""" + return litellm.Usage(**payload) + + def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: for key in keys: value = getattr(details, key, None) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cc1fa1b8803..6d700bae514 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7844,3 +7844,102 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream(): async with _wrapper_over(stream): pass assert stream.close_calls == 1 + + +def _final_chunk_with_dict_usage(choices) -> ModelResponseStream: + """A stream chunk whose ``usage`` is a plain dict. + + ``ModelResponseStream.__init__`` coerces a dict ``usage`` kwarg into + ``litellm.Usage``, so the dict is attached after construction — the same + shape a custom ``streaming_decoder`` (or any third-party iterator that + bypasses litellm chunk construction) delivers to ``chunk_creator``. + """ + chunk: Final = ModelResponseStream( + id="chatcmpl-dict-usage-test", + created=1754900000, + model="gpt-4o", + object="chat.completion.chunk", + choices=choices, + ) + chunk.usage = { # mutable-ok: wire-shaped usage payload under test + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "prompt_tokens_details": {"cached_tokens": 80}, # mutable-ok: wire payload + "completion_tokens_details": {"reasoning_tokens": 5}, # mutable-ok: wire payload + "cache_creation_input_tokens": 42, + } + return chunk + + +def test_dispatch_provider_chunk_dict_usage_preserves_token_details( + logging_obj: Logging, +): + """The openai else-branch of ``_dispatch_provider_chunk`` handles three + usage shapes: ``Usage`` (passed through), ``BaseModel`` + (``Usage(**model_dump())``), and plain ``dict``. The dict arm used to + rebuild ``litellm.Usage`` from only prompt/completion/total tokens, + silently dropping ``prompt_tokens_details`` (cached_tokens), + ``completion_tokens_details`` (reasoning_tokens), and cache token counts + that both sibling arms preserve. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage( + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=None)) + ] # mutable-ok: chunk payload + ) + model_response: Final = wrapper.model_response_creator() + wrapper._dispatch_provider_chunk( + chunk=chunk, + model_response=model_response, + completion_obj={"content": ""}, # mutable-ok: dispatch scratch dict + ) + + usage: Final = model_response.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 10 + assert usage.total_tokens == 110 + assert usage.prompt_tokens_details is not None, ( + "dict-usage arm dropped prompt_tokens_details (its Usage/BaseModel sibling arms preserve it)" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 + assert getattr(usage, "cache_creation_input_tokens", None) == 42 + + +def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details( + logging_obj: Logging, +): + """End-to-end through ``chunk_creator``: a usage-only final chunk (empty + ``choices``, ``stream_options.include_usage`` set) takes the early-return + path, so whatever the dict-usage arm wrote is exactly what the consumer + receives — there is no later aggregation to repair it. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage(choices=[]) # mutable-ok: empty choices list + result: Final = wrapper.chunk_creator(chunk=chunk) + + assert result is not None + usage: Final = result.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens_details is not None, ( + "usage-only early return delivered a Usage stripped of prompt_tokens_details to the stream consumer" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 From 9b2fc9db40f8f976597388b4450ff002ee0aa732 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 27 Oct 2026 15:26:15 -0700 Subject: [PATCH 027/147] test: remove checks that pin implementation details and add a shadcn testing rule (#48663) * test: remove e2e checks and a CI line count that pin implementation details Deletes the Auto Router template select anchoring spec, the 390px logs pagination variant, and the call id tooltip and clipboard step in the logs spec. Each one tests library or test helper behavior rather than what users see, and each has been red or flaky on main without a product regression Also drops the exact line count on the fan-in requirements export in the integration coverage job. It went from 6 to 10 when uv.lock gained a Python 3.15 fork, while the install was still correct. A package the filter drops still fails the job, either at the hash-checked install or at the next step's import * docs(agents): do not test what shadcn and Base UI already guarantee Adds the boundary to the dashboard AGENTS.md next to the test tiers, and points the Playwright suite's entry in tests/e2e/AGENTS.md at it --- .circleci/config.yml | 1 - tests/e2e/AGENTS.md | 2 +- tests/e2e/ui/tests/logs/logs.spec.ts | 11 --- .../e2e/ui/tests/logs/logsPagination.spec.ts | 75 +++++++++---------- .../autoRouterTemplateSelect.spec.ts | 52 ------------- ui/litellm-dashboard/AGENTS.md | 2 + 6 files changed, 38 insertions(+), 105 deletions(-) delete mode 100644 tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts diff --git a/.circleci/config.yml b/.circleci/config.yml index 1edf22578ff..4ce5ea8b6fe 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1379,7 +1379,6 @@ jobs: uv export --frozen --all-groups --all-extras --no-emit-project --no-emit-workspace --no-header --no-annotate \ | awk '/^[^ ]/ { keep = $0 ~ /^(annotated-types|coverage|pydantic|pydantic-core|typing-extensions|typing-inspection)==/ } keep' \ > /tmp/fan-in-requirements.txt - test "$(grep -c '^[^ ]' /tmp/fan-in-requirements.txt)" -eq 6 uv venv .venv --python 3.12 uv pip install --no-config --python .venv --require-hashes -r /tmp/fan-in-requirements.txt - attach_workspace: diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index ef57135a4d2..cf1e41a59a6 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -49,7 +49,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke -- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` +- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json`. Its specs follow the dashboard's rule against testing what shadcn and Base UI already guarantee, in `ui/litellm-dashboard/AGENTS.md` ## MCP suite: real Datadog only diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index 688bfe8e3b4..49400f1509a 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -194,17 +194,6 @@ test.describe("Logs page", () => { await expect(row, `no logs row for call id ${callId}`).toHaveCount(1, { timeout: 30_000 }); await expect(row, "the row itself shows only the request id").not.toContainText(callId); - await row.getByText(requestId).hover(); - const tooltip = page.locator("[data-slot='tooltip-content']"); - await expect(tooltip, "hovering the Request ID cell does not list the x-litellm-call-id").toContainText( - `x-litellm-call-id: ${callId}`, - { timeout: 10_000 }, - ); - await tooltip.getByRole("button", { name: "Copy x-litellm-call-id" }).click(); - if (await page.evaluate(() => window.isSecureContext)) { - await expect.poll(() => page.evaluate(() => navigator.clipboard.readText())).toBe(callId); - } - await row.click(); const drawer = page.getByRole("dialog").first(); await expect(drawer.getByText("Request & Response")).toBeVisible({ timeout: 20_000 }); diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index 838a829a3e2..585eb31c37c 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -54,52 +54,47 @@ async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): test.describe("Logs pagination with Aria in the header", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - for (const viewport of [ - { width: 1440, height: 900 }, - { width: 390, height: 844 }, - ]) { - test(`keeps pagination reachable and Aria available at ${viewport.width}px`, async ({ page, request }) => { - await page.setViewportSize(viewport); - const alias = `e2e-aria-header-${uniqueSuffix()}`; - const key = await createVirtualKey(request, { key_alias: alias }); - try { - const ids = await Promise.all( - Array.from({ length: 26 }, (_, index) => - sendChatCompletion(request, { - model: CHAT_MODEL_A, - prompt: `${alias}-${index}`, - apiKey: key.key, - }), - ), - ); - await Promise.all(ids.map((id) => waitForSpendLog(request, id))); - await openLogs(page); - await setAriaHidden(page, false); - await applyKeyAliasFilter(page, await openFilterDrawer(page), alias); - await setRowsPerPage(page, "25"); - await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); - await visibleTestId(page, "pagination-next").click(); - await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26"); + test("keeps pagination reachable and Aria available at 1440px", async ({ page, request }) => { + await page.setViewportSize({ width: 1440, height: 900 }); + const alias = `e2e-aria-header-${uniqueSuffix()}`; + const key = await createVirtualKey(request, { key_alias: alias }); + try { + const ids = await Promise.all( + Array.from({ length: 26 }, (_, index) => + sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `${alias}-${index}`, + apiKey: key.key, + }), + ), + ); + await Promise.all(ids.map((id) => waitForSpendLog(request, id))); + await openLogs(page); + await setAriaHidden(page, false); + await applyKeyAliasFilter(page, await openFilterDrawer(page), alias); + await setRowsPerPage(page, "25"); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + await visibleTestId(page, "pagination-next").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26"); - const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + await aria.click(); + await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); + await page.keyboard.press("Escape"); + await expect(aria).toBeFocused(); + await visibleTestId(page, "pagination-first").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { + await navigateToPage(page, destination); await aria.click(); await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); await page.keyboard.press("Escape"); await expect(aria).toBeFocused(); - await visibleTestId(page, "pagination-first").click(); - await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); - for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { - await navigateToPage(page, destination); - await aria.click(); - await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); - await page.keyboard.press("Escape"); - await expect(aria).toBeFocused(); - } - } finally { - await deleteVirtualKey(request, key.token); } - }); - } + } finally { + await deleteVirtualKey(request, key.token); + } + }); }); test.describe("Logs page session-grouped pagination", () => { diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts deleted file mode 100644 index 5f05953cc80..00000000000 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ /dev/null @@ -1,52 +0,0 @@ -import { expect, test, type Locator, type Page as PlaywrightPage } from "@playwright/test"; -import { ADMIN_STORAGE_PATH } from "../../constants"; -import { navigateToPage } from "../../helpers/navigation"; -import { Page } from "../../fixtures/pages"; - -/** - * Opens Add Auto Router and returns the Template select's trigger, which is the - * shallowest real page that renders SelectContent with tall multi-line options. - */ -async function openTemplateSelect(page: PlaywrightPage) { - await navigateToPage(page, Page.Models); - await page.getByRole("tab", { name: "Auto-Routers" }).click(); - await page.getByRole("button", { name: "Add Auto Router" }).click(); - - const trigger = page.getByTestId("template-selector"); - await expect(trigger).toBeVisible(); - return trigger; -} - -async function boxes(trigger: Locator, options: Locator) { - const triggerBox = await trigger.boundingBox(); - const optionsBox = await options.boundingBox(); - return triggerBox && optionsBox ? { triggerBox, optionsBox } : null; -} - -const clippedPopup = (page: PlaywrightPage) => page.locator('[data-slot="select-content"]'); - -function pollOptionsCoverTrigger(trigger: Locator, options: Locator) { - return expect.poll(async () => { - const box = await boxes(trigger, options); - return ( - box && - box.optionsBox.y < box.triggerBox.y + box.triggerBox.height && - box.optionsBox.y + box.optionsBox.height > box.triggerBox.y - ); - }); -} - -test.describe("Auto Router template select anchoring", () => { - test.use({ storageState: ADMIN_STORAGE_PATH }); - - test("keeps the trigger uncovered when the options open with no room below it", async ({ page }) => { - await page.setViewportSize({ width: 1280, height: 560 }); - const trigger = await openTemplateSelect(page); - await trigger.scrollIntoViewIfNeeded(); - - await trigger.click(); - await expect(page.getByRole("listbox")).toBeVisible(); - - await pollOptionsCoverTrigger(trigger, clippedPopup(page)).toBe(false); - }); -}); diff --git a/ui/litellm-dashboard/AGENTS.md b/ui/litellm-dashboard/AGENTS.md index b6fe5219022..f5576b1228b 100644 --- a/ui/litellm-dashboard/AGENTS.md +++ b/ui/litellm-dashboard/AGENTS.md @@ -8,6 +8,8 @@ When you fix lint violations that are grandfathered in `eslint-suppressions.json Tests come in three tiers, named by the standard definitions. `Foo.test.tsx` is a unit test: one module, collaborators replaced by doubles, no multi-component tree, and it should run in milliseconds. `Foo.integration.test.tsx` renders a real component tree with real children and only stubs the network boundary. It costs seconds per case, so it earns its place by proving wiring that a unit test cannot reach. Browser-level tests live in `tests/e2e/ui/` as Playwright specs against a live proxy +Do not test what shadcn and Base UI already guarantee, such as a tooltip opening on hover, a popup positioning itself against its trigger, focus trapping and Escape in a dialog, keyboard movement through a menu or select, and open and close animations. Those are the library's contract, and a test of them only fails when we restyle or move the page. Test what we put inside the primitive and what happens because of it, like the text a tooltip carries, the options a select offers, or the request a confirm button sends. Reach that content the most direct way, by a click, the keyboard, or the data the page renders, rather than through a hover race. If we change a primitive's behavior in `src/components/ui/`, its test lives next to that primitive as `*.test.tsx`, never in a page spec that borrows the page as a host + When a component holds logic worth asserting, extract the logic and unit-test it there rather than driving it through a render. `CreateMCPServer` is the worked example: its payload building lives in `createServerPayload.ts` with 46 unit tests that run in single-digit milliseconds, while `CreateMCPServer.integration.test.tsx` keeps only the cases that prove a form field reaches the right payload key. A test that renders a whole modal to assert the shape of one object belongs in the first category, not the second Most of the suite predates this split and is not yet classified, so an unsuffixed `*.test.tsx` is not evidence that a file is really a unit test. Classify what you touch From c4c2fe6ba12c9b76f43e8bf3cde5ce63efb91a2b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 18:38:07 -0400 Subject: [PATCH 028/147] docs(router): clarify heuristic v2 keyword settings (#48666) --- .../complexity_router/README.md | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 08d39841b61..3df1803d264 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -13,7 +13,7 @@ Unlike the semantic `auto_router` which uses embedding-based matching, the `comp ## How It Works -The router scores each request across 7 dimensions: +The default `classifier_type: heuristic` scores each request across 7 dimensions. These weights, keyword lists and score boundaries apply to the weighted heuristic scorer. [`heuristic_v2`](#heuristic-v2) uses a separate success-probability predictor | Dimension | Description | Weight | |-----------|-------------|--------| @@ -254,11 +254,23 @@ model_list: No classifier model call or per-model training data is required. The classifier uses global tier quality, request-type quality, and similar-request cohorts from -the bundled UltraFeedback artifact. It estimates success at every tier, enforces +the bundled UltraFeedback artifact. Cohorts group prompts by request type, length, +code and math markers, multiple-choice formatting and non-ASCII character share. +Each tier's global estimate is `(successes + 1) / (observations + 2)`. Request-type +and cohort estimates combine their successes and observations with the broader +estimate as a prior, so sparse or missing statistics fall back toward broader data. +The predictor estimates success at every tier, enforces monotonic probabilities, and returns the first tier meeting the success threshold, or REASONING if no tier meets it. The existing complexity-router tier pool then selects and dispatches a model from that tier +`heuristic_v2` does not read `reasoning_keywords` or apply `reasoning_override_min_score`, +even when the floor is `0`. The weighted scorer's keyword lists, `dimension_weights`, +`token_thresholds` and `tier_boundaries` do not affect v2's probabilities or tier selection. +This also applies when a chain selects `local_heuristic: heuristic_v2`. Use +`classifier_type: heuristic` for the [reasoning-marker override](#reasoning-override). +The separate `keyword_tier_rules` routing overrides still run before classification + Set `heuristic_v2_success_threshold` to a value from 0 to 1 to override the artifact's threshold. For example, `0.9` requires a predicted success probability of at least 90%. Higher thresholds favor more capable tiers. Omit the setting or @@ -646,7 +658,16 @@ of that threshold, or when no tier meets the threshold. A zero margin still defe ### Reasoning Override -If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone. +The weighted heuristic scorer promotes its result to REASONING when the user message matches 2+ entries in `reasoning_keywords` and the weighted score reaches `reasoning_override_min_score`. The floor tracks `tier_boundaries.simple_medium` unless set. An explicit `0` lowers the floor to zero, so a negative weighted score still does not qualify. `heuristic_v2` bypasses this scorer and does not apply the override + +For routing based on these reasoning keywords, set the following under your existing `complexity_router_config` and keep your tier mapping + +```yaml +classifier_type: heuristic +reasoning_override_min_score: 0 +``` + +Omit `reasoning_keywords` to use the built-in list, or supply a nonempty list to replace it. These scoring keywords are separate from `keyword_tier_rules`, which can select a tier before either classifier runs ### System Prompt Handling From 969161501acfca278642ea645130754e66dee846 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 27 Oct 2026 15:54:39 -0700 Subject: [PATCH 029/147] feat(ui): add LITELLM_UI_USE_SI_TERMINOLOGY to show SI instead of AI in the Admin UI (#47440) * feat(ui): add LITELLM_UI_USE_SI_TERMINOLOGY to show SI instead of AI in the Admin UI Setting LITELLM_UI_USE_SI_TERMINOLOGY=true on the proxy makes the Admin UI say "SI" where its own copy says "AI" and "Super Intelligence" where it says "Artificial Intelligence". This covers the nav, page headings, settings text, guardrail garden blurbs, the compliance playground prompts and the playground sample prompts. Third-party product, provider and law names keep their spelling (Vertex AI, Cisco AI Defense, EU AI Act, ...), and user data such as guardrail names, gateway URLs and keys is never rewritten. The proxy exposes the flag on /.well-known/litellm-ui-config, a TerminologyProvider at the app root reads it, and components call useTerminology() on their own strings. The context defaults to a pass-through, so nothing changes when the flag is unset. * fix(ui): apply SI terminology to the command palette, page visibility groups and guardrail garden search * test(ui): cover the flag-off wording on the palette, page visibility and garden search tests * fix(ui): apply SI terminology to the Aria usage assistant copy * test(ui): keep the Aria ui-config stub within the inline object lint budget * fix(ui): keep AI wording in prompts sent to models --- .../ui_discovery_endpoints.py | 2 + .../ui_discovery_endpoints.py | 1 + .../test_ui_discovery_endpoints.py | 22 ++++++ .../agents/_components/AgentsPanel.tsx | 7 +- .../_components/EvaluationSettingsModal.tsx | 6 +- .../_components/guardrail_garden.test.tsx | 30 +++++++- .../_components/guardrail_garden.tsx | 6 +- .../_components/guardrail_garden_card.tsx | 4 +- .../_components/guardrail_garden_detail.tsx | 6 +- .../_components/MCPGatewaySessionsTab.tsx | 17 +++-- .../mcp-servers/_components/mcp_connect.tsx | 12 ++-- .../components/chat_ui/RealtimePlayground.tsx | 8 ++- .../components/complianceUI/ComplianceUI.tsx | 4 +- .../_components/ai_suggestion_modal.tsx | 4 +- .../_components/guardrail_selection_modal.tsx | 4 +- .../policies/_components/policy_templates.tsx | 4 +- .../_components/template_parameter_modal.tsx | 4 +- .../UsageViewSelect/UsageViewSelect.tsx | 4 +- .../components/overview/TopAgents.tsx | 4 +- ui/litellm-dashboard/src/app/layout.tsx | 9 ++- .../src/components/AIHub/ModelHubTable.tsx | 4 +- .../CommandPalette.integration.test.tsx | 28 ++++++++ .../CommandPalette/CommandPalette.tsx | 4 +- .../components/CommandPalette/utils.test.ts | 11 +++ .../src/components/CommandPalette/utils.ts | 7 +- .../src/components/DashboardHeader.tsx | 4 +- .../Navbar/ViewSwitcher.integration.test.tsx | 24 +++++-- .../src/components/Navbar/ViewSwitcher.tsx | 9 ++- .../PageVisibilitySettings.test.tsx | 32 ++++++++- .../UISettings/PageVisibilitySettings.tsx | 8 ++- .../AdminSettings/UISettings/UISettings.tsx | 12 +++- .../src/components/leftnav.test.tsx | 19 ++++- .../src/components/leftnav.tsx | 15 ++-- .../liteadmin/LiteAdmin.integration.test.tsx | 69 ++++++++++++++----- .../src/components/liteadmin/LiteAdmin.tsx | 7 +- .../components/model_add/CredentialsPanel.tsx | 4 +- .../components/model_add/CredentialsTable.tsx | 4 +- .../src/components/moyai/MoyaiConnected.tsx | 4 +- .../src/components/networking.tsx | 1 + .../organisms/create_key_button.tsx | 11 ++- .../shared/CreatedKeyDisplay.test.tsx | 8 +++ .../components/shared/CreatedKeyDisplay.tsx | 9 ++- .../templates/KeyEditViewControls.tsx | 51 ++++++++------ .../src/contexts/TerminologyContext.test.tsx | 37 ++++++++++ .../src/contexts/TerminologyContext.tsx | 18 +++++ .../src/data/compliancePrompts.test.ts | 33 +++++++++ .../src/data/compliancePrompts.ts | 11 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 ++ .../src/lib/terminology.test.ts | 45 ++++++++++++ ui/litellm-dashboard/src/lib/terminology.ts | 28 ++++++++ 50 files changed, 567 insertions(+), 113 deletions(-) create mode 100644 ui/litellm-dashboard/src/contexts/TerminologyContext.test.tsx create mode 100644 ui/litellm-dashboard/src/contexts/TerminologyContext.tsx create mode 100644 ui/litellm-dashboard/src/data/compliancePrompts.test.ts create mode 100644 ui/litellm-dashboard/src/lib/terminology.test.ts create mode 100644 ui/litellm-dashboard/src/lib/terminology.ts diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 8225984a964..792ef8a3f8c 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -25,6 +25,7 @@ async def get_ui_config(): or general_settings.get("auto_redirect_ui_login_to_sso", False) is True ) admin_ui_disabled: Final = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" + use_si_terminology: Final = os.getenv("LITELLM_UI_USE_SI_TERMINOLOGY", "false").lower() == "true" hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings) sso_configured: Final = has_user_setup_sso() @@ -49,4 +50,5 @@ async def get_ui_config(): is_control_plane=is_control_plane, workers=proxy_config.worker_registry if is_control_plane else [], mcp_stdio_enabled=is_mcp_stdio_enabled(), + use_si_terminology=use_si_terminology, ) diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index e30ad26bfe5..d396c79b844 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -14,4 +14,5 @@ class UiDiscoveryEndpoints(LiteLLMBaseModel): hide_default_credentials_hint: bool = False is_control_plane: bool = False mcp_stdio_enabled: bool = False + use_si_terminology: bool = False workers: list[WorkerRegistryEntry] = Field(default=[]) diff --git a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index 1ef0eee761a..d338f38d862 100644 --- a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -501,3 +501,25 @@ def test_ui_config_tells_the_dashboard_whether_stdio_mcp_servers_are_enabled(mon assert response.status_code == 200 assert response.json()["mcp_stdio_enabled"] is expected + + +@pytest.mark.parametrize(("flag", "expected"), [(None, False), ("false", False), ("true", True), ("TRUE", True)]) +def test_ui_config_tells_the_dashboard_whether_to_use_si_terminology( + monkeypatch: pytest.MonkeyPatch, flag: str | None, expected: bool +) -> None: + if flag is None: + monkeypatch.delenv("LITELLM_UI_USE_SI_TERMINOLOGY", raising=False) + else: + monkeypatch.setenv("LITELLM_UI_USE_SI_TERMINOLOGY", flag) + app: Final[FastAPI] = FastAPI() + app.include_router(router) + + with ( + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), + ): + response: Final = TestClient(app).get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + assert response.json()["use_si_terminology"] is expected diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx index b392e270d33..3f599c38976 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect } from "react"; import { Info, Plus } from "lucide-react"; import { getAgentsList, deleteAgentCall } from "@/components/networking"; @@ -31,6 +32,7 @@ interface AgentsResponse { } const AgentsPanel: React.FC = ({ accessToken, userRole, teams }) => { + const term = useTerminology(); const [agentsList, setAgentsList] = useState([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(true); @@ -140,8 +142,9 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams

Agents

- List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents - public. + {term( + "List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents public.", + )}

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx index 900a04e480d..78091584cde 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import { Play } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; @@ -47,6 +48,7 @@ export function EvaluationSettingsModal({ accessToken, onRunEvaluation, }: EvaluationSettingsModalProps) { + const term = useTerminology(); const [prompt, setPrompt] = useState(DEFAULT_PROMPT); const [schema, setSchema] = useState(DEFAULT_SCHEMA); const [model, setModel] = useState(null); @@ -95,8 +97,8 @@ export function EvaluationSettingsModal({ Evaluation Settings {guardrailName - ? `Configure AI evaluation for ${guardrailName}` - : "Configure AI evaluation for re-running on logs"} + ? `${term("Configure AI evaluation for")} ${guardrailName}` + : term("Configure AI evaluation for re-running on logs")} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx index 613fee04f68..ee214b451fa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx @@ -1,9 +1,16 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { render, screen, within } from "@testing-library/react"; +import { fireEvent, render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import { TerminologyProvider } from "@/contexts/TerminologyContext"; import GuardrailGarden from "./guardrail_garden"; import { ALL_CARDS } from "./guardrail_garden_data"; +const uiConfig = vi.hoisted(() => ({ data: undefined as { use_si_terminology?: boolean } | undefined })); + +vi.mock("@/app/(dashboard)/hooks/uiConfig/useUIConfig", () => ({ + useUIConfig: () => ({ data: uiConfig.data }), +})); + vi.mock("./guardrail_garden_detail", () => ({ __esModule: true, default: ({ card, onBack }: { card: { name: string }; onBack: () => void }) => ( @@ -21,6 +28,7 @@ const DECISION_MODEL_CARDS = ALL_CARDS.filter((c) => c.category === "decision_mo describe("GuardrailGarden", () => { beforeEach(() => { vi.clearAllMocks(); + uiConfig.data = undefined; }); const renderGarden = () => render(); @@ -114,6 +122,26 @@ describe("GuardrailGarden", () => { } }); + it.each([ + [true, "SI"], + [false, "AI"], + ])("should match the wording a card shows when SI terminology is %s", (useSi, shownTerm) => { + uiConfig.data = { use_si_terminology: useSi }; + const card = ALL_CARDS.find((c) => c.description.startsWith("AI security platform protecting")); + expect(card).toBeDefined(); + render( + + + , + ); + + fireEvent.change(screen.getByPlaceholderText("Search guardrails"), { + target: { value: `${shownTerm} security platform protecting` }, + }); + + expect(screen.getByText(card!.name)).toBeInTheDocument(); + }); + it("should show an empty result set for a query that matches nothing", async () => { const user = userEvent.setup(); renderGarden(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx index f8e1ab1ef66..d430a5a9f19 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState } from "react"; import { ArrowRight, Search } from "lucide-react"; import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; @@ -11,6 +12,7 @@ interface GuardrailGardenProps { } const GuardrailGarden: React.FC = ({ accessToken, onGuardrailCreated }) => { + const term = useTerminology(); const [searchQuery, setSearchQuery] = useState(""); const [selectedCard, setSelectedCard] = useState(null); const [showAllLitellm, setShowAllLitellm] = useState(false); @@ -22,7 +24,7 @@ const GuardrailGarden: React.FC = ({ accessToken, onGuardr const q = searchQuery.toLowerCase(); return ( card.name.toLowerCase().includes(q) || - card.description.toLowerCase().includes(q) || + term(card.description).toLowerCase().includes(q) || card.tags.some((t) => t.toLowerCase().includes(q)) ); }); @@ -106,7 +108,7 @@ const GuardrailGarden: React.FC = ({ accessToken, onGuardr Partner Guardrails

- Third-party guardrail integrations from leading AI security providers. + {term("Third-party guardrail integrations from leading AI security providers.")}

{partnerCards.map((card) => ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx index 2b7cedd6d85..8a6375d0d68 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx @@ -1,9 +1,11 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React from "react"; import { CircleCheck } from "lucide-react"; import { GuardrailCardInfo } from "./guardrail_garden_data"; import { Logo } from "@/components/molecules/logo/Logo"; const GuardrailCard: React.FC<{ card: GuardrailCardInfo; onClick: () => void }> = ({ card, onClick }) => { + const term = useTerminology(); return (
void }> {card.name}
-

{card.description}

+

{term(card.description)}

{card.eval && (
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx index 8ac0493b012..e549ebca50a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState } from "react"; import { ArrowLeft } from "lucide-react"; import { Button } from "@/components/ui/button"; @@ -21,6 +22,7 @@ interface GuardrailDetailViewProps { } const GuardrailDetailView: React.FC = ({ card, onBack, accessToken, onGuardrailCreated }) => { + const term = useTerminology(); const [isAddFormVisible, setIsAddFormVisible] = useState(false); const [activeTab, setActiveTab] = useState("overview"); @@ -66,7 +68,7 @@ const GuardrailDetailView: React.FC = ({ card, onBack,

{card.name}

-

{card.description}

+

{term(card.description)}

{/* Action buttons — outlined style like Vertex */}
@@ -101,7 +103,7 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {/* Left column — overview + details table */}

Overview

-

{card.description}

+

{term(card.description)}

Guardrail Details

Details are as follows

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx index 04a095f8792..46c1d2c6848 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx @@ -1,5 +1,6 @@ "use client"; +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState } from "react"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { RefreshCw, Unplug } from "lucide-react"; @@ -132,6 +133,7 @@ function SessionsBody({ isLoading: boolean; onDisconnect: ((selector: MCPGatewaySessionSelector) => void) | null; }) { + const term = useTerminology(); if (isLoading) { return (

- No live MCP connections on this worker (pid {data.worker_pid}). Connect an AI client to the gateway to see it - here. + No live MCP connections on this worker (pid {data.worker_pid}).{" "} + {term("Connect an AI client to the gateway to see it here.")}

); @@ -166,11 +168,11 @@ function SessionsBody({ <>
- +
- + (null); const queryOptions = { @@ -282,9 +285,9 @@ export function MCPGatewaySessionsTab({ accessToken, canTerminate }: MCPGatewayS

Live Connections

- Stateful Streamable HTTP sessions currently open on this proxy worker, grouped by the AI client that sent - the MCP initialize request and by the authenticated LiteLLM user. Stateless requests and SSE connections are - not counted. + {term( + "Stateful Streamable HTTP sessions currently open on this proxy worker, grouped by the AI client that sent the MCP initialize request and by the authenticated LiteLLM user. Stateless requests and SSE connections are not counted.", + )}

- Use tools directly from Cursor IDE with LiteLLM MCP. Enable your AI assistant to perform real-world tasks - without leaving your coding environment. + {term( + "Use tools directly from Cursor IDE with LiteLLM MCP. Enable your AI assistant to perform real-world tasks without leaving your coding environment.", + )}
@@ -475,8 +478,9 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = []

Connect to your MCP client

- Use tools directly from any MCP client with LiteLLM MCP. Enable your AI assistant to perform real-world - tasks through a simple, secure connection. + {term( + "Use tools directly from any MCP client with LiteLLM MCP. Enable your AI assistant to perform real-world tasks through a simple, secure connection.", + )}

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx index 8acc54bee8d..01151203e65 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx @@ -1,5 +1,6 @@ "use client"; +import { useTerminology } from "@/contexts/TerminologyContext"; import { CircleX, Mic, MicOff, Send, Volume2 } from "lucide-react"; import React, { useCallback, useEffect, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; @@ -27,6 +28,7 @@ const RealtimePlayground: React.FC = ({ customProxyBaseUrl, selectedGuardrails, }) => { + const term = useTerminology(); const [messages, setMessages] = useState([]); const [inputText, setInputText] = useState(""); const [isConnected, setIsConnected] = useState(false); @@ -408,8 +410,8 @@ const RealtimePlayground: React.FC = ({ Realtime Voice Playground

- Click Connect to start a realtime session. You can speak using your microphone or type messages. - The AI will respond with voice and text. + Click Connect to start a realtime session. You can speak using your microphone or type messages.{" "} + {term("The AI will respond with voice and text.")}

)} @@ -428,7 +430,7 @@ const RealtimePlayground: React.FC = ({ : "bg-muted text-foreground rounded-bl-md" }`} > -
{msg.role === "user" ? "You" : "AI"}
+
{msg.role === "user" ? "You" : term("AI")}
{msg.content}
)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index 54ea3702cb4..db576ed03a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -7,6 +7,7 @@ import { type CompliancePrompt, } from "@/data/compliancePrompts"; import useCan from "@/app/(dashboard)/hooks/useCan"; +import { useTerminology } from "@/contexts/TerminologyContext"; import { getGuardrailsList, testPoliciesAndGuardrails } from "@/components/networking"; import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector"; import { Policy } from "@/components/policies/types"; @@ -131,7 +132,8 @@ export default function ComplianceUI({ proxySettings, }: ComplianceUIProps) { const canViewPolicies = useCan("viewPolicies"); - const frameworks = getFrameworks(); + const term = useTerminology(); + const frameworks = getFrameworks(term); const [policyValueToLabel, setPolicyValueToLabel] = useState>(new Map()); const [guardrailOptions, setGuardrailOptions] = useState([]); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx index c4f76b07a4e..50ffb89ae83 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useEffect, useMemo, useState } from "react"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; @@ -62,6 +63,7 @@ const AiSuggestionModal: React.FC = ({ accessToken, allTemplates, }) => { + const term = useTerminology(); const [attackExamples, setAttackExamples] = useState([""]); const [description, setDescription] = useState(""); const [isLoading, setIsLoading] = useState(false); @@ -815,7 +817,7 @@ const AiSuggestionModal: React.FC = ({ {/* Header */}
- AI Policy Suggestion + {term("AI Policy Suggestion")}

{showResults ? `${suggestions?.length || 0} template${(suggestions?.length || 0) !== 1 ? "s" : ""} matched your requirements` diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx index c4a45d88d1f..25a4b89ed2a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect } from "react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -40,6 +41,7 @@ const GuardrailSelectionModal: React.FC = ({ isLoading = false, progressInfo, }) => { + const term = useTerminology(); const [selectedGuardrails, setSelectedGuardrails] = useState>(new Set()); // Prepare guardrail info with existence status @@ -199,7 +201,7 @@ const GuardrailSelectionModal: React.FC = ({

✨ - AI-Discovered Competitors ({template.discoveredCompetitors.length}) + {term("AI-Discovered Competitors")} ({template.discoveredCompetitors.length})
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx index 46dda1dd11c..0db1f74f540 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect, useMemo } from "react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -106,6 +107,7 @@ const PolicyTemplates: React.FC = ({ onTemplatesLoaded, accessToken, }) => { + const term = useTerminology(); const [templates, setTemplates] = useState([]); const [isLoading, setIsLoading] = useState(false); const [selectedTags, setSelectedTags] = useState>(new Set()); @@ -191,7 +193,7 @@ const PolicyTemplates: React.FC = ({ - Use AI to find templates + {term("Use AI to find templates")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx index b90379b149d..6aa26088756 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect } from "react"; import { Button } from "@/components/ui/button"; import { Badge } from "@/components/ui/badge"; @@ -44,6 +45,7 @@ const TemplateParameterModal: React.FC = ({ isLoading = false, accessToken, }) => { + const term = useTerminology(); const [parameterValues, setParameterValues] = useState>({}); const [competitorMode, setCompetitorMode] = useState<"ai" | "manual">("ai"); const [selectedModel, setSelectedModel] = useState(null); @@ -260,7 +262,7 @@ const TemplateParameterModal: React.FC = ({ className="grid-cols-2" >