From 29691d6dddce8d9d4cbe53370b0b00be1975da6d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:39:15 -0700 Subject: [PATCH] fix(spend-logs): match short error codes on the indexed prefix alone A code under 128 characters is its own prefix, so the prefix comparison is already exact and Postgres no longer reads each matching row's metadata to recheck the full code. Codes of 128 characters or more keep the full-code check. The integration test now seeds long codes that share a prefix, and error_code is declared with Annotated so direct calls default to None. --- .../spend_management_endpoints.py | 21 +++++++--- .../spend/test_spend_log_error_filters.py | 10 +++++ .../test_spend_management_endpoints.py | 42 ++++++++++++++++++- 3 files changed, 67 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index b98624ff998..09fc2285f51 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -96,6 +96,9 @@ _SESSION_REPRESENTATIVE_ORDER_SQL: Final = ( f'call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC, request_id' ) _BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_batch', 'retrieve_batch')" +_ERROR_CODE_SQL: Final = "metadata->'error_information'->>'error_code'" +_INDEXED_ERROR_CODE_PREFIX_LENGTH: Final = 128 +_INDEXED_ERROR_CODE_PREFIX_SQL: Final = f"left({_ERROR_CODE_SQL}, {_INDEXED_ERROR_CODE_PREFIX_LENGTH})" _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType( { "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}", @@ -226,6 +229,15 @@ class _SessionSpendRow(TypedDict): session_models: ReadOnly[Sequence[str]] +def _error_code_condition(placeholder: str, error_code: str) -> str: + if len(error_code) < _INDEXED_ERROR_CODE_PREFIX_LENGTH: + return f"{_INDEXED_ERROR_CODE_PREFIX_SQL} = {placeholder}" + return ( + f"{_INDEXED_ERROR_CODE_PREFIX_SQL} = left({placeholder}, {_INDEXED_ERROR_CODE_PREFIX_LENGTH}) " + f"AND {_ERROR_CODE_SQL} = {placeholder}" + ) + + _SESSION_MODELS_LIMIT: Final = 10 _SESSION_MODEL_NAME_MAX_LEN: Final = 256 @@ -2577,7 +2589,9 @@ async def ui_view_spend_logs( model_group: str | None = fastapi.Query(default=None, description="Filter logs by model group"), key_alias: str | None = fastapi.Query(default=None, description="Filter logs by key alias"), end_user: str | None = fastapi.Query(default=None, description="Filter logs by end user"), - error_code: str | None = fastapi.Query(default=None, description="Filter logs by error code (e.g., '404', '500')"), + error_code: Annotated[ + str | None, fastapi.Query(description="Filter logs by error code (e.g., '404', '500')") + ] = None, error_message: str | None = fastapi.Query( default=None, description="Filter logs by error message (partial string match)" ), @@ -2933,10 +2947,7 @@ async def ui_view_spend_logs( sql_params.append(f"%{key_alias}%") p += 1 if error_code is not None: - sql_conditions.append( - f"left(metadata->'error_information'->>'error_code', 128) = left(${p}, 128) " - f"AND metadata->'error_information'->>'error_code' = ${p}" - ) + sql_conditions.append(_error_code_condition(f"${p}", error_code)) sql_params.append(error_code) p += 1 if error_message is not None: diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py index 13bd1405fd1..a1e3b97e5db 100644 --- a/tests/integration/spend/test_spend_log_error_filters.py +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -1,3 +1,4 @@ +import secrets import time from collections.abc import Mapping from dataclasses import dataclass @@ -22,6 +23,10 @@ INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" SUCCESS_ROWS: Final = 5000 STATS_FLUSH_WINDOW_SECONDS: Final = 11.0 +LONG_CODE_PREFIX: Final = secrets.token_hex(64) +LONG_CODE: Final = LONG_CODE_PREFIX + secrets.token_hex(1500) +LONG_CODE_SIBLING: Final = LONG_CODE_PREFIX + secrets.token_hex(1500) + @dataclass(frozen=True, slots=True) class ErrorRow: @@ -41,11 +46,16 @@ ERROR_ROWS: Final = ( ErrorRow("failed-500-timeout", "failure", "500", "upstream timed out"), ErrorRow("failed-429", "failure", "429", "rate limited by upstream"), ErrorRow("client-disconnected", "success", "499", "Client disconnected the request"), + ErrorRow("failed-long-code", "failure", LONG_CODE, "provider error with a long code"), + ErrorRow("failed-long-code-sibling", "failure", LONG_CODE_SIBLING, "provider error with a long code"), ) ERROR_CODES: Final = MappingProxyType( { "500": frozenset({"failed-500-overloaded", "failed-500-timeout"}), "499": frozenset({"client-disconnected"}), + LONG_CODE: frozenset({"failed-long-code"}), + LONG_CODE_PREFIX: frozenset(), + LONG_CODE_PREFIX[:-1]: frozenset(), } ) SESSION_VIEWS: Final = ("false", "true") diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 8d53817f51e..ca08863330c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -124,7 +124,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): gte = re.search(r'"startTime" >= \(\$(\d+)', cond) lte = re.search(r'"startTime" <= \(\$(\d+)', cond) alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond) - code = re.search(r"error_code' = \$(\d+)", cond) + code = re.search(r"error_code'(?:, 128\))? = \$(\d+)", cond) msg = re.search(r"error_message' LIKE \$(\d+)", cond) credential = re.fullmatch(r"metadata->>'used_client_oauth_token' = \$(\d+)", cond) sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond) @@ -4770,6 +4770,46 @@ async def test_ui_view_spend_logs_with_error_code(client): app.dependency_overrides.pop(ps.user_api_key_auth, None) + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client): + prefix = "E" * 128 + mock_spend_logs = [ + { + "id": f"log-{suffix}", + "request_id": f"req-{suffix}", + "api_key": "sk-test-key", + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + "metadata": json.dumps({"error_information": {"error_code": prefix + suffix}}), + } + for suffix in ("first", "second") + ] + + def filter_by_exact_error_code(where): + wanted = str(where.get("metadata", {}).get("equals", "")).strip('"') + return [log for log in mock_spend_logs if json.loads(log["metadata"])["error_information"]["error_code"] == wanted] + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + with patch.object( + ps, "prisma_client", make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_exact_error_code) + ): + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"error_code": prefix + "first", "start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert [row["id"] for row in response.json()["data"]] == ["log-first"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + @pytest.mark.asyncio async def test_ui_view_spend_logs_with_error_message(client): """Test filtering spend logs by error message"""