mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_spend_log_request_id_call_id
# Conflicts: # litellm/proxy/spend_tracking/spend_management_endpoints.py
This commit is contained in:
mateo-berri
commit
0c2be34fb1
28 changed files with 1242 additions and 189 deletions
|
|
@ -121,6 +121,23 @@ class TokenEndpointClient:
|
|||
return Ok(ExchangedToken(access_token=parsed.access_token, expires_in=parsed.expires_in))
|
||||
|
||||
|
||||
class _KeyGuard:
|
||||
"""The per-key single-flight lock plus the invalidation generation that lock protects.
|
||||
|
||||
Both live on one object so their lifetimes cannot diverge. `get_or_compute` binds the guard to
|
||||
a local for its whole critical section, which keeps the weak map's entry alive for as long as
|
||||
that compute could still write; an `invalidate` overlapping the compute therefore reaches the
|
||||
very same object and its bump is guaranteed to be observed. Conversely a guard nobody holds is
|
||||
collectible precisely because no write is outstanding for it to fence.
|
||||
"""
|
||||
|
||||
__slots__ = ("__weakref__", "generation", "lock")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.lock = asyncio.Lock()
|
||||
self.generation = 0
|
||||
|
||||
|
||||
class ExchangedTokenCache:
|
||||
"""Memoizes the final token string per key, single-flighting concurrent misses on one lock."""
|
||||
|
||||
|
|
@ -129,7 +146,7 @@ class ExchangedTokenCache:
|
|||
max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
|
||||
default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
|
||||
)
|
||||
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
|
||||
self._guards: weakref.WeakValueDictionary[str, _KeyGuard] = weakref.WeakValueDictionary()
|
||||
|
||||
async def get_or_compute(
|
||||
self,
|
||||
|
|
@ -144,28 +161,50 @@ class ExchangedTokenCache:
|
|||
guaranteeing the token it gets back was minted for the *current* inputs: a stored entry
|
||||
whose fingerprint differs reads as a miss and is re-minted over. That keeps eviction
|
||||
addressable without the key having to encode the credential material it protects.
|
||||
|
||||
An `invalidate` landing while `compute` is in flight wins over that compute's write. The
|
||||
token is still returned to the caller it was minted for, but it is not stored, so the next
|
||||
resolution re-mints rather than serving a bearer that predates the invalidation for the
|
||||
rest of its TTL.
|
||||
"""
|
||||
cached = self._get(cache_key, fingerprint)
|
||||
if cached is not None:
|
||||
return Ok(cached)
|
||||
async with self._lock(cache_key):
|
||||
guard = self._guard(cache_key)
|
||||
async with guard.lock:
|
||||
cached = self._get(cache_key, fingerprint)
|
||||
if cached is not None:
|
||||
return Ok(cached)
|
||||
generation = guard.generation
|
||||
match await compute():
|
||||
case Ok(token):
|
||||
self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
cache_key,
|
||||
(fingerprint, token.access_token),
|
||||
ttl=_cache_ttl_seconds(token.expires_in),
|
||||
)
|
||||
if guard.generation == generation:
|
||||
self._store(cache_key, fingerprint, token)
|
||||
return Ok(token.access_token)
|
||||
case Error(err):
|
||||
return Error(err)
|
||||
|
||||
def invalidate(self, cache_key: str) -> None:
|
||||
"""Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401)."""
|
||||
"""Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401).
|
||||
|
||||
Bumping the guard's generation is what makes the eviction stick against a compute already
|
||||
awaiting the token endpoint: that compute snapshotted the old generation and so skips its
|
||||
write. No guard means no compute is in flight, since an in-flight one pins its own.
|
||||
|
||||
Stays synchronous: callers invalidate from plain `def`s.
|
||||
"""
|
||||
self._cache.delete_cache(cache_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
guard = self._guards.get(cache_key)
|
||||
if guard is None:
|
||||
return
|
||||
guard.generation += 1
|
||||
|
||||
def _store(self, cache_key: str, fingerprint: str, token: ExchangedToken) -> None:
|
||||
self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
cache_key,
|
||||
(fingerprint, token.access_token),
|
||||
ttl=_cache_ttl_seconds(token.expires_in),
|
||||
)
|
||||
|
||||
def _get(self, cache_key: str, fingerprint: str) -> str | None:
|
||||
"""The stored token, or None when absent or minted for different inputs.
|
||||
|
|
@ -180,12 +219,12 @@ class ExchangedTokenCache:
|
|||
return None
|
||||
return token if stored_fingerprint == fingerprint else None
|
||||
|
||||
def _lock(self, cache_key: str) -> asyncio.Lock:
|
||||
lock = self._locks.get(cache_key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
self._locks[cache_key] = lock
|
||||
return lock
|
||||
def _guard(self, cache_key: str) -> _KeyGuard:
|
||||
guard = self._guards.get(cache_key)
|
||||
if guard is None:
|
||||
guard = _KeyGuard()
|
||||
self._guards[cache_key] = guard
|
||||
return guard
|
||||
|
||||
|
||||
def _cache_ttl_seconds(expires_in: int | None) -> int:
|
||||
|
|
|
|||
|
|
@ -3855,6 +3855,13 @@ if MCP_AVAILABLE:
|
|||
and server.auth_type == MCPAuth.oauth2_token_exchange
|
||||
and oauth2_headers
|
||||
and len(mcp_servers or []) == 1
|
||||
and server.server_id
|
||||
in frozenset(
|
||||
allowed.server_id
|
||||
for allowed in await _get_allowed_mcp_servers(
|
||||
user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip
|
||||
)
|
||||
)
|
||||
):
|
||||
await global_mcp_server_manager.preflight_token_exchange(
|
||||
server=server,
|
||||
|
|
|
|||
|
|
@ -100,9 +100,10 @@ class CliPollData(TypedDict, total=False):
|
|||
|
||||
|
||||
class CliSsoStartData(TypedDict):
|
||||
login_id: str
|
||||
poll_secret: str
|
||||
user_code: str
|
||||
login_id: ReadOnly[str]
|
||||
poll_secret: ReadOnly[str]
|
||||
user_code: ReadOnly[str]
|
||||
verification_uri_complete: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
class CliAuthResult(TypedDict):
|
||||
|
|
@ -860,11 +861,22 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
|
|||
poll_secret: Final = cli_sso_flow["poll_secret"]
|
||||
user_code: Final = cli_sso_flow["user_code"]
|
||||
|
||||
sso_url = f"{base_url}/sso/key/generate?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id})
|
||||
browser_prefills_code: Final = isinstance(cli_sso_flow.get("verification_uri_complete"), str)
|
||||
sso_url: Final = f"{base_url}/sso/key/generate?" + urlencode(
|
||||
(
|
||||
("source", LITELLM_CLI_SOURCE_IDENTIFIER),
|
||||
("key", key_id),
|
||||
*((("user_code", user_code),) if browser_prefills_code else ()),
|
||||
)
|
||||
)
|
||||
|
||||
click.echo(f"Opening browser to: {sso_url}")
|
||||
click.echo("Please complete the SSO authentication in your browser...")
|
||||
click.echo(f"Verification code: {user_code}")
|
||||
click.echo(
|
||||
f"Verification code: {user_code} (pre-filled in the browser, check it matches)"
|
||||
if browser_prefills_code
|
||||
else f"Verification code: {user_code}"
|
||||
)
|
||||
click.echo(f"Session ID: {key_id}")
|
||||
|
||||
# Open browser
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .crowdstrike_aidr import CrowdStrikeAIDRHandler
|
||||
from .crowdstrike_aidr import CrowdStrikeAIDRHandler, streaming_params_from_litellm_params
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
|
@ -15,17 +15,16 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
if not guardrail_name:
|
||||
raise ValueError("CrowdStrike AIDR guardrail name is required")
|
||||
|
||||
streaming_params: Final = streaming_params_from_litellm_params(litellm_params)
|
||||
_crowdstrike_aidr_callback: Final = CrowdStrikeAIDRHandler(
|
||||
guardrail_name=guardrail_name,
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
# Exclude during_call to prevent duplicate input events
|
||||
event_hook=[
|
||||
GuardrailEventHooks.pre_call.value,
|
||||
GuardrailEventHooks.post_call.value,
|
||||
],
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
fail_on_error=litellm_params.fail_on_error,
|
||||
streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
|
||||
streaming_sampling_rate=streaming_params.streaming_sampling_rate,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(_crowdstrike_aidr_callback)
|
||||
|
||||
|
|
|
|||
|
|
@ -24,8 +24,11 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
|
||||
from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionToolParam
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import (
|
||||
CrowdStrikeAIDRGuardrailConfigModelOptionalParams,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -153,6 +156,21 @@ def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Mapping[str, Any] |
|
|||
return merged if present else None
|
||||
|
||||
|
||||
def streaming_params_from_litellm_params(
|
||||
litellm_params: LitellmParams,
|
||||
) -> CrowdStrikeAIDRGuardrailConfigModelOptionalParams:
|
||||
extras: Final[Mapping[str, object]] = litellm_params.model_extra or {}
|
||||
nested: Final = litellm_params.optional_params
|
||||
optional_params: Final[Mapping[str, object]] = {} if nested is None else nested.model_dump()
|
||||
return CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_validate(
|
||||
{
|
||||
name: value
|
||||
for name in CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_fields
|
||||
if (value := optional_params.get(name, extras.get(name))) is not None
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _messages_since_last_assistant(
|
||||
messages: Sequence[AllMessageValues],
|
||||
) -> _FilteredMessages:
|
||||
|
|
@ -241,6 +259,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
fail_on_error: bool | None = True,
|
||||
streaming_end_of_stream_only: bool | None = None,
|
||||
streaming_sampling_rate: int | None = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -250,10 +270,19 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
|
|||
guardrail_name (str): The name of the guardrail instance.
|
||||
api_key (str | None): The CrowdStrike AIDR API key. Reads from CS_AIDR_TOKEN env var if None.
|
||||
api_base (str | None): The CrowdStrike AIDR API base URL. Reads from CS_AIDR_BASE_URL env var if None.
|
||||
streaming_end_of_stream_only (bool | None): Scan streamed output once at end of stream instead of
|
||||
every streaming_sampling_rate chunks. Defaults to False.
|
||||
streaming_sampling_rate (int | None): Scan the accumulated streamed output every Nth chunk. Defaults to 5.
|
||||
**kwargs: Additional arguments passed to the CustomGuardrail base class.
|
||||
"""
|
||||
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
self.fail_on_error = True if fail_on_error is None else fail_on_error
|
||||
self._set_streaming_params(
|
||||
CrowdStrikeAIDRGuardrailConfigModelOptionalParams(
|
||||
streaming_end_of_stream_only=streaming_end_of_stream_only,
|
||||
streaming_sampling_rate=streaming_sampling_rate,
|
||||
)
|
||||
)
|
||||
|
||||
self.api_key = api_key or os.environ.get("CS_AIDR_TOKEN")
|
||||
if not self.api_key:
|
||||
|
|
@ -274,6 +303,15 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
|
|||
"Initialized CrowdStrike AIDR Guardrail: name=%s, api_base=%s", guardrail_name, self.api_base
|
||||
)
|
||||
|
||||
def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None:
|
||||
self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False
|
||||
self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5
|
||||
|
||||
@override
|
||||
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
|
||||
super().update_in_memory_litellm_params(litellm_params)
|
||||
self._set_streaming_params(streaming_params_from_litellm_params(litellm_params))
|
||||
|
||||
async def _call_crowdstrike_aidr_guard(
|
||||
self, payload: dict[str, Any], hook_name: str
|
||||
) -> _GuardChatCompletionsResult:
|
||||
|
|
|
|||
|
|
@ -56,6 +56,10 @@ router: Final = APIRouter()
|
|||
|
||||
SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000
|
||||
|
||||
_SESSION_GROUP_KEY_SQL: Final = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
|
||||
_MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')"
|
||||
_AGENT_CALL_TYPE_SQL: Final = "'asend_message'"
|
||||
|
||||
_INTERNAL_HEALTH_CHECK_API_KEYS: Final = (
|
||||
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
hash_token(token=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME),
|
||||
|
|
@ -145,21 +149,16 @@ class _DailyTagSpendRow(TypedDict):
|
|||
total_spend: float
|
||||
|
||||
|
||||
class _SessionCountAggregate(TypedDict):
|
||||
session_id: int
|
||||
|
||||
|
||||
class _SessionCountRow(TypedDict):
|
||||
session_id: str
|
||||
_count: _SessionCountAggregate
|
||||
|
||||
|
||||
class _SessionSpendRow(TypedDict):
|
||||
session_id: str
|
||||
api_key: ReadOnly[str]
|
||||
session_total_count: ReadOnly[int]
|
||||
session_total_spend: float
|
||||
mcp_tool_call_count: int
|
||||
mcp_tool_call_spend: float
|
||||
session_cache_hit_count: ReadOnly[int]
|
||||
session_llm_count: ReadOnly[int]
|
||||
session_agent_count: ReadOnly[int]
|
||||
|
||||
|
||||
class _SpendSumAggregate(TypedDict, total=False):
|
||||
|
|
@ -268,18 +267,6 @@ async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, obj
|
|||
return await _spend_logs_table(prisma_client).count(where=where)
|
||||
|
||||
|
||||
async def _count_logs_per_session(
|
||||
prisma_client: PrismaClient, session_ids: Sequence[str | None]
|
||||
) -> Sequence[_SessionCountRow]:
|
||||
"""Count spend log rows per session for the given session ids."""
|
||||
rows: Final = await _spend_logs_table(prisma_client).group_by(
|
||||
by=["session_id"],
|
||||
where={"session_id": {"in": session_ids}},
|
||||
count={"session_id": True},
|
||||
)
|
||||
return cast(Sequence[_SessionCountRow], rows) # cast-ok: group_by(count=) shape is fixed by the by/count args
|
||||
|
||||
|
||||
async def _find_team_row(prisma_client: PrismaClient, team_id: str) -> _SupportsModelDump | None:
|
||||
"""Read a single team row as a Prisma model instance."""
|
||||
return await _team_table(prisma_client).find_unique(where={"team_id": team_id})
|
||||
|
|
@ -2316,6 +2303,10 @@ async def ui_view_spend_logs(
|
|||
default=False,
|
||||
description="Exclude LiteLLM internal health check requests from results",
|
||||
),
|
||||
group_by_session: bool = fastapi.Query(
|
||||
default=False,
|
||||
description="Paginate over sessions instead of raw logs: one representative row per session, total counts sessions",
|
||||
),
|
||||
):
|
||||
"""
|
||||
View spend logs with pagination support.
|
||||
|
|
@ -2674,12 +2665,16 @@ async def ui_view_spend_logs(
|
|||
else:
|
||||
_order_expr = order_column
|
||||
|
||||
joined_conditions: Final = " AND ".join(sql_conditions)
|
||||
session_grouping: Final = group_by_session is True
|
||||
count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else ""
|
||||
count_query: Final = f"""
|
||||
SELECT COUNT(*) AS total_count
|
||||
FROM (
|
||||
SELECT 1
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE {" AND ".join(sql_conditions)}
|
||||
WHERE {joined_conditions}
|
||||
{count_group_clause}
|
||||
LIMIT ${p}
|
||||
) AS bounded_matches
|
||||
"""
|
||||
|
|
@ -2690,9 +2685,7 @@ async def ui_view_spend_logs(
|
|||
total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total
|
||||
|
||||
sql_query: Final = f"""
|
||||
SELECT
|
||||
request_id, call_type, api_key, spend, total_tokens,
|
||||
select_columns: Final = """request_id, call_type, api_key, spend, total_tokens,
|
||||
prompt_tokens, completion_tokens, "startTime", "endTime",
|
||||
"completionStartTime", model, model_id, model_group,
|
||||
custom_llm_provider, api_base, "user", metadata,
|
||||
|
|
@ -2700,12 +2693,29 @@ async def ui_view_spend_logs(
|
|||
organization_id, end_user, requester_ip_address,
|
||||
session_id, status, mcp_namespaced_tool_name, agent_id,
|
||||
litellm_call_id,
|
||||
COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms
|
||||
COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms"""
|
||||
sql_query: Final = (
|
||||
f"""
|
||||
SELECT * FROM (
|
||||
SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
|
||||
{select_columns}
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE {joined_conditions}
|
||||
ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
|
||||
) AS session_representatives
|
||||
ORDER BY {exact_request_id_first}{_order_expr} {_sql_dir}{_nulls_clause}, request_id
|
||||
LIMIT ${p} OFFSET ${p + 1}
|
||||
"""
|
||||
if session_grouping
|
||||
else f"""
|
||||
SELECT
|
||||
{select_columns}
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE {" AND ".join(sql_conditions)}
|
||||
WHERE {joined_conditions}
|
||||
ORDER BY {exact_request_id_first}{_order_expr} {_sql_dir}{_nulls_clause}
|
||||
LIMIT ${p} OFFSET ${p + 1}
|
||||
"""
|
||||
)
|
||||
sql_params.extend([page_size, skip])
|
||||
|
||||
data: Final = await prisma_client.db.query_raw(sql_query, *sql_params)
|
||||
|
|
@ -4122,11 +4132,12 @@ async def _build_ui_spend_logs_response(
|
|||
Build the paginated response for the UI spend-logs endpoint.
|
||||
|
||||
When ``enrich_session_counts`` is ``True`` (the default for the v1/UI
|
||||
endpoint), each row is enriched with ``session_total_count`` so the
|
||||
frontend knows which sessions are expandable (multi-call sessions).
|
||||
For every row that carries a ``session_id``, a single ``GROUP BY`` query
|
||||
fetches the total number of logs in each referenced session. Rows without
|
||||
a ``session_id`` default to ``1``.
|
||||
endpoint), each row is enriched with ``session_total_count`` plus spend
|
||||
and call-type aggregates so the frontend knows which sessions are
|
||||
expandable (multi-call sessions). One ``GROUP BY (session_id, api_key)``
|
||||
query serves every referenced session, keyed per api key so two callers
|
||||
reusing a session id never see each other's totals. Rows without a
|
||||
``session_id`` default to ``1``.
|
||||
|
||||
When ``enrich_session_counts`` is ``False`` (v2 endpoint), rows are
|
||||
serialised without the extra query.
|
||||
|
|
@ -4148,7 +4159,6 @@ async def _build_ui_spend_logs_response(
|
|||
A dict with ``data`` (enriched rows), ``total``, ``page``,
|
||||
``page_size``, ``total_pages``, and ``total_is_capped``.
|
||||
"""
|
||||
count_map: dict[str, int] = {}
|
||||
if enrich_session_counts:
|
||||
session_ids: Final[Sequence[str | None]] = list(
|
||||
{
|
||||
|
|
@ -4157,15 +4167,8 @@ async def _build_ui_spend_logs_response(
|
|||
if (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None))
|
||||
}
|
||||
)
|
||||
if session_ids:
|
||||
# NOTE: This GROUP BY runs on every v1/UI page load. The IN clause
|
||||
# is bounded by page_size (typically 25-50 distinct session IDs).
|
||||
# If performance degrades at scale, consider short-lived caching or
|
||||
# folding the count into the main query via a window function.
|
||||
counts: Final = await _count_logs_per_session(prisma_client, session_ids)
|
||||
count_map = {r["session_id"]: r["_count"]["session_id"] for r in counts if r.get("session_id")}
|
||||
|
||||
session_spend_map: dict[str, dict[str, int | float]] = {}
|
||||
session_spend_map: dict[tuple[str, str], dict[str, int | float]] = {}
|
||||
if enrich_session_counts and session_ids:
|
||||
from prisma.errors import PrismaError
|
||||
|
||||
|
|
@ -4177,38 +4180,46 @@ async def _build_ui_spend_logs_response(
|
|||
{
|
||||
(row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None))
|
||||
for row in data
|
||||
if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None))
|
||||
if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) is not None
|
||||
}
|
||||
)
|
||||
rows: Final[Sequence[_SessionSpendRow]] = await _query_raw(
|
||||
prisma_client,
|
||||
"""
|
||||
SELECT session_id,
|
||||
f"""
|
||||
SELECT session_id, api_key,
|
||||
COUNT(*)::int AS session_total_count,
|
||||
COALESCE(SUM(spend), 0)::double precision AS session_total_spend,
|
||||
COUNT(*) FILTER (
|
||||
WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools')
|
||||
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
|
||||
)::int AS mcp_tool_call_count,
|
||||
COALESCE(SUM(spend) FILTER (
|
||||
WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools')
|
||||
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
|
||||
), 0)::double precision AS mcp_tool_call_spend,
|
||||
COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count
|
||||
COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count,
|
||||
COUNT(*) FILTER (
|
||||
WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
|
||||
)::int AS session_llm_count,
|
||||
COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE session_id = ANY($1::text[])
|
||||
AND api_key = ANY($2::text[])
|
||||
GROUP BY session_id
|
||||
GROUP BY session_id, api_key
|
||||
""",
|
||||
session_ids,
|
||||
authorized_api_keys,
|
||||
)
|
||||
session_spend_map = {
|
||||
row["session_id"]: {
|
||||
(row["session_id"], row["api_key"]): {
|
||||
"session_total_count": int(row.get("session_total_count") or 0),
|
||||
"session_total_spend": float(row.get("session_total_spend") or 0.0),
|
||||
"mcp_tool_call_count": int(row.get("mcp_tool_call_count") or 0),
|
||||
"mcp_tool_call_spend": float(row.get("mcp_tool_call_spend") or 0.0),
|
||||
"session_cache_hit_count": int(row.get("session_cache_hit_count") or 0),
|
||||
"session_llm_count": int(row.get("session_llm_count") or 0),
|
||||
"session_agent_count": int(row.get("session_agent_count") or 0),
|
||||
}
|
||||
for row in rows
|
||||
if row.get("session_id")
|
||||
if row.get("session_id") and row.get("api_key") is not None
|
||||
}
|
||||
except PrismaError:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -4221,14 +4232,17 @@ async def _build_ui_spend_logs_response(
|
|||
for row in data:
|
||||
row_dict = dict(row) if isinstance(row, dict) else row.model_dump()
|
||||
sid = row_dict.get("session_id")
|
||||
row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1
|
||||
session_stats = session_spend_map.get(sid) if sid else None
|
||||
row_api_key = row_dict.get("api_key")
|
||||
session_stats = session_spend_map.get((sid, row_api_key)) if sid and row_api_key is not None else None
|
||||
row_dict["session_total_count"] = int(session_stats["session_total_count"]) if session_stats else 1
|
||||
if session_stats:
|
||||
row_dict["session_total_spend"] = session_stats["session_total_spend"]
|
||||
if session_stats["mcp_tool_call_count"]:
|
||||
row_dict["mcp_tool_call_count"] = session_stats["mcp_tool_call_count"]
|
||||
row_dict["mcp_tool_call_spend"] = session_stats["mcp_tool_call_spend"]
|
||||
row_dict["session_cache_hit_count"] = session_stats["session_cache_hit_count"]
|
||||
row_dict["session_llm_count"] = session_stats["session_llm_count"]
|
||||
row_dict["session_agent_count"] = session_stats["session_agent_count"]
|
||||
enriched.append(row_dict)
|
||||
response_data: list = enriched
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from typing_extensions import TypeIs
|
|||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
EMPTY_MAPPING,
|
||||
LITELLM_MAX_STREAMING_DURATION_SECONDS,
|
||||
STREAM_SSE_DONE_STRING,
|
||||
)
|
||||
|
|
@ -273,6 +274,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._hidden_params["additional_headers"] = process_response_headers(
|
||||
self.response.headers or {}
|
||||
) # GUARANTEE OPENAI HEADERS IN RESPONSE
|
||||
self._raw_response_headers: Mapping[str, str] = MappingProxyType(
|
||||
dict(self.response.headers or {}) # mutable-ok: immediately frozen by MappingProxyType
|
||||
)
|
||||
|
||||
def _check_max_streaming_duration(self) -> None:
|
||||
"""Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS."""
|
||||
|
|
@ -446,6 +450,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
except Exception:
|
||||
# Fallback to original if serialization fails
|
||||
pass
|
||||
self._restore_provider_response_headers(logging_response)
|
||||
|
||||
end_time: Final = datetime.now()
|
||||
if is_async:
|
||||
|
|
@ -480,6 +485,41 @@ class BaseResponsesAPIStreamingIterator:
|
|||
)
|
||||
self._run_post_success_hooks(end_time=end_time)
|
||||
|
||||
def _restore_provider_response_headers(self, logging_response: object) -> None:
|
||||
"""Re-apply the provider's response headers to the copy handed to logging callbacks.
|
||||
|
||||
``model_validate(model_dump())`` above drops pydantic private attributes, so the
|
||||
``_hidden_params`` the provider transform set on the nested response are lost. Returns early
|
||||
when that copy fell back to the original event, so logging-only state never lands on the
|
||||
object the caller is iterating.
|
||||
"""
|
||||
if logging_response is self.completed_response:
|
||||
return
|
||||
target: Final[object] = getattr(logging_response, "response", None)
|
||||
existing_hidden: Final[object] = getattr(target, "_hidden_params", None)
|
||||
if not isinstance(existing_hidden, Mapping):
|
||||
return
|
||||
existing: Final[Mapping[str, object]] = existing_hidden
|
||||
source_hidden: Final[object] = getattr(
|
||||
getattr(self.completed_response, "response", None), "_hidden_params", None
|
||||
)
|
||||
source: Final[Mapping[str, object]] = source_hidden if isinstance(source_hidden, Mapping) else EMPTY_MAPPING
|
||||
processed: Final[object] = source.get("additional_headers") or self._hidden_params.get("additional_headers")
|
||||
raw: Final[object] = source.get("headers") or self._raw_response_headers
|
||||
headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING
|
||||
raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING
|
||||
# rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy
|
||||
# splats into the client's HTTP headers, and copying non-header keys would carry response_cost
|
||||
setattr( # noqa: B010 # target is typed object here, so a plain attribute store does not type check
|
||||
target,
|
||||
"_hidden_params",
|
||||
{ # mutable-ok: the cost calculator writes optional_params into _hidden_params
|
||||
"additional_headers": {**headers}, # mutable-ok: fresh copy, logging callbacks may mutate it
|
||||
"headers": {**raw_headers}, # mutable-ok: fresh copy, logging callbacks may mutate it
|
||||
**existing,
|
||||
},
|
||||
)
|
||||
|
||||
def _handle_logging_completed_response(self):
|
||||
"""Base implementation - should be overridden by subclasses"""
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,18 @@ from .base import GuardrailConfigModel
|
|||
|
||||
|
||||
class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel):
|
||||
pass
|
||||
streaming_end_of_stream_only: bool | None = Field(
|
||||
default=None,
|
||||
description="If False (default when unset), post_call scans the accumulated streamed response every "
|
||||
"streaming_sampling_rate chunks and an in-flight block stops the stream. If True, the guard runs once "
|
||||
"over the assembled response at end of stream, so flagged content may already have reached the client.",
|
||||
)
|
||||
streaming_sampling_rate: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description="When streaming_end_of_stream_only is False, scan the accumulated streamed response every Nth "
|
||||
"chunk. Defaults to 5 when unset.",
|
||||
)
|
||||
|
||||
|
||||
class CrowdStrikeAIDRGuardrailConfigModel(GuardrailConfigModel[CrowdStrikeAIDRGuardrailConfigModelOptionalParams]):
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ interface ChatOptions {
|
|||
apiKey?: string;
|
||||
/** Sent as `user`, which lands in the spend log's end_user column. */
|
||||
endUser?: string;
|
||||
/** Sent as `litellm_trace_id`, which lands in the spend log's session_id column. */
|
||||
traceId?: string;
|
||||
}
|
||||
|
||||
/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */
|
||||
|
|
@ -34,6 +36,7 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO
|
|||
model: opts.model,
|
||||
messages: [{ role: "user", content: opts.prompt }],
|
||||
...(opts.endUser ? { user: opts.endUser } : {}),
|
||||
...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}),
|
||||
},
|
||||
});
|
||||
expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true);
|
||||
|
|
|
|||
150
tests/e2e/ui/tests/logs/logsPagination.spec.ts
Normal file
150
tests/e2e/ui/tests/logs/logsPagination.spec.ts
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
|
||||
import { ADMIN_STORAGE_PATH } from "../../constants";
|
||||
import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
|
||||
import { Page } from "../../fixtures/pages";
|
||||
import { CHAT_MODEL_A, createVirtualKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
|
||||
|
||||
/**
|
||||
* Session-grouped pagination (#38060): a page of N rows must render exactly N session rows, a
|
||||
* session must never straddle pages, and two callers reusing one session id stay separate rows.
|
||||
* All traffic is generated per run behind a unique key alias or session id, so concurrent specs
|
||||
* cannot decide the outcome.
|
||||
*/
|
||||
|
||||
const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
||||
|
||||
/** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */
|
||||
const requestLogsRows = (page: PlaywrightPage): Locator =>
|
||||
page.locator("table").filter({ visible: true }).first().locator("tbody tr");
|
||||
|
||||
const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
|
||||
|
||||
async function openLogs(page: PlaywrightPage): Promise<void> {
|
||||
await navigateToPage(page, Page.Logs);
|
||||
await dismissFeedbackPopup(page);
|
||||
await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 });
|
||||
}
|
||||
|
||||
async function openFilterDrawer(page: PlaywrightPage): Promise<Locator> {
|
||||
await visibleTestId(page, "datatable-filters-trigger").click();
|
||||
const drawer = page.getByRole("dialog", { name: "Filters" });
|
||||
await expect(drawer).toBeVisible({ timeout: 10_000 });
|
||||
return drawer;
|
||||
}
|
||||
|
||||
async function applyKeyAliasFilter(page: PlaywrightPage, drawer: Locator, alias: string): Promise<void> {
|
||||
await drawer.getByRole("combobox", { name: "Search a key alias" }).click();
|
||||
await page.keyboard.type(alias);
|
||||
await page.getByRole("option", { name: alias, exact: true }).first().click();
|
||||
await drawer.getByRole("button", { name: "Apply Filters" }).click();
|
||||
await expect(drawer).not.toBeVisible({ timeout: 10_000 });
|
||||
}
|
||||
|
||||
async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise<void> {
|
||||
await visibleTestId(page, "pagination-page-size").click();
|
||||
await page.getByRole("option", { name: size, exact: true }).click();
|
||||
}
|
||||
|
||||
test.describe("Logs page session-grouped pagination", () => {
|
||||
test.use({ storageState: ADMIN_STORAGE_PATH });
|
||||
|
||||
test("a 25-row page renders exactly 25 session rows and no session straddles pages", async ({ page, request }) => {
|
||||
const suffix = uniqueSuffix();
|
||||
const alias = `e2e-logs-pgn-${suffix}`;
|
||||
const mine = await createVirtualKey(request, { key_alias: alias });
|
||||
|
||||
const soloIds: string[] = [];
|
||||
for (let i = 0; i < 26; i++) {
|
||||
soloIds.push(
|
||||
await sendChatCompletion(request, {
|
||||
model: CHAT_MODEL_A,
|
||||
prompt: `logs-pgn-solo-${i}-${suffix}`,
|
||||
apiKey: mine.key,
|
||||
}),
|
||||
);
|
||||
}
|
||||
const sessionA = `sess-pgn-a-${suffix}`;
|
||||
const sessionB = `sess-pgn-b-${suffix}`;
|
||||
let lastSessionCallId = "";
|
||||
for (let i = 0; i < 7; i++) {
|
||||
lastSessionCallId = await sendChatCompletion(request, {
|
||||
model: CHAT_MODEL_A,
|
||||
prompt: `logs-pgn-a-${i}-${suffix}`,
|
||||
apiKey: mine.key,
|
||||
traceId: sessionA,
|
||||
});
|
||||
}
|
||||
for (let i = 0; i < 3; i++) {
|
||||
lastSessionCallId = await sendChatCompletion(request, {
|
||||
model: CHAT_MODEL_A,
|
||||
prompt: `logs-pgn-b-${i}-${suffix}`,
|
||||
apiKey: mine.key,
|
||||
traceId: sessionB,
|
||||
});
|
||||
}
|
||||
await waitForSpendLog(request, lastSessionCallId);
|
||||
await waitForSpendLog(request, soloIds[soloIds.length - 1]);
|
||||
|
||||
// 36 calls in 28 session groups: 26 solos plus sessions of 7 and 3.
|
||||
await openLogs(page);
|
||||
const drawer = await openFilterDrawer(page);
|
||||
await applyKeyAliasFilter(page, drawer, alias);
|
||||
await setRowsPerPage(page, "25");
|
||||
|
||||
await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 28", { timeout: 30_000 });
|
||||
await expect(requestLogsRows(page)).toHaveCount(25);
|
||||
// The sessions are the newest groups, so their single representative rows sit on page 1.
|
||||
await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(1);
|
||||
await expect(requestLogsRows(page).filter({ hasText: sessionA })).toContainText("7");
|
||||
await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(1);
|
||||
|
||||
await visibleTestId(page, "pagination-next").click();
|
||||
|
||||
await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-28 of 28", { timeout: 30_000 });
|
||||
await expect(requestLogsRows(page)).toHaveCount(3);
|
||||
await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(0);
|
||||
await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(0);
|
||||
});
|
||||
|
||||
test("two keys reusing one session id stay separate rows", async ({ page, request }) => {
|
||||
const suffix = uniqueSuffix();
|
||||
const mine = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-mine-${suffix}` });
|
||||
const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-theirs-${suffix}` });
|
||||
const sharedSession = `sess-pgn-shared-${suffix}`;
|
||||
|
||||
let lastId = "";
|
||||
for (let i = 0; i < 2; i++) {
|
||||
lastId = await sendChatCompletion(request, {
|
||||
model: CHAT_MODEL_A,
|
||||
prompt: `logs-pgn-shared-mine-${i}-${suffix}`,
|
||||
apiKey: mine.key,
|
||||
traceId: sharedSession,
|
||||
});
|
||||
}
|
||||
lastId = await sendChatCompletion(request, {
|
||||
model: CHAT_MODEL_A,
|
||||
prompt: `logs-pgn-shared-theirs-${suffix}`,
|
||||
apiKey: theirs.key,
|
||||
traceId: sharedSession,
|
||||
});
|
||||
await waitForSpendLog(request, lastId);
|
||||
|
||||
await openLogs(page);
|
||||
const drawer = await openFilterDrawer(page);
|
||||
await drawer.getByPlaceholder("Enter session ID…").fill(sharedSession);
|
||||
await drawer.getByRole("button", { name: "Apply Filters" }).click();
|
||||
await expect(drawer).not.toBeVisible({ timeout: 10_000 });
|
||||
|
||||
// One row per caller: reusing a session id must not merge two keys' activity into one row.
|
||||
await expect(requestLogsRows(page).filter({ hasText: sharedSession })).toHaveCount(2, { timeout: 30_000 });
|
||||
|
||||
// And each row carries ITS key's totals: two calls badge the first key's row,
|
||||
// while the other key's single call renders as a plain LLM row.
|
||||
const mineRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: mine.token });
|
||||
const theirsRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: theirs.token });
|
||||
await expect(mineRow).toHaveCount(1);
|
||||
await expect(theirsRow).toHaveCount(1);
|
||||
await expect(mineRow.getByText("2", { exact: true })).toBeVisible();
|
||||
await expect(theirsRow.getByText("LLM", { exact: true })).toBeVisible();
|
||||
});
|
||||
});
|
||||
|
|
@ -7,6 +7,7 @@ cache's hit/single-flight behavior. Each assertion fails under a real mutation o
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -359,6 +360,110 @@ async def test_cache_invalidate_only_evicts_the_named_key():
|
|||
assert calls == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_invalidate_mid_compute_is_not_overwritten_by_that_compute():
|
||||
"""A bearer minted before an invalidation must never be served after it.
|
||||
|
||||
The compute is suspended at the token endpoint when the invalidation lands, so its write is
|
||||
the one that would resurrect the evicted bearer for the rest of its TTL. The caller it was
|
||||
minted for still gets it; the *cache* is what the invalidation is about.
|
||||
"""
|
||||
cache = ExchangedTokenCache()
|
||||
mint_started, release_mint = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def slow_mint():
|
||||
mint_started.set()
|
||||
await release_mint.wait()
|
||||
return _ok_token("bearer-minted-before-invalidation")
|
||||
|
||||
async def re_mint():
|
||||
return _ok_token("bearer-minted-after-invalidation")
|
||||
|
||||
in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp"))
|
||||
await mint_started.wait()
|
||||
|
||||
assert not in_flight.done()
|
||||
cache.invalidate("slot")
|
||||
release_mint.set()
|
||||
|
||||
raced = await in_flight
|
||||
assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation"
|
||||
|
||||
after = await cache.get_or_compute("slot", re_mint, fingerprint="fp")
|
||||
assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_invalidate_mid_compute_survives_garbage_collection():
|
||||
"""The record of an invalidation must outlive a collection cycle taken mid-compute.
|
||||
|
||||
Per-key state is held weakly so idle keys do not accumulate. If the state a compute checks
|
||||
before writing were collectible while that compute is suspended, the check would read as
|
||||
"nothing was invalidated" and the stale write would land; the running compute has to pin it.
|
||||
"""
|
||||
cache = ExchangedTokenCache()
|
||||
mint_started, release_mint = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def slow_mint():
|
||||
mint_started.set()
|
||||
await release_mint.wait()
|
||||
return _ok_token("bearer-minted-before-invalidation")
|
||||
|
||||
async def re_mint():
|
||||
return _ok_token("bearer-minted-after-invalidation")
|
||||
|
||||
in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp"))
|
||||
await mint_started.wait()
|
||||
|
||||
assert not in_flight.done()
|
||||
cache.invalidate("slot")
|
||||
gc.collect()
|
||||
release_mint.set()
|
||||
await in_flight
|
||||
|
||||
after = await cache.get_or_compute("slot", re_mint, fingerprint="fp")
|
||||
assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_stores_a_compute_that_started_after_the_invalidation():
|
||||
"""Only the mint that predates the invalidation loses its write.
|
||||
|
||||
A caller queued behind the single-flight lock computes after the eviction, so its token is
|
||||
fresh and must be cached; otherwise the fix would trade one stale bearer for re-minting on
|
||||
every subsequent resolution.
|
||||
"""
|
||||
cache = ExchangedTokenCache()
|
||||
mint_started, release_mint = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def slow_mint():
|
||||
mint_started.set()
|
||||
await release_mint.wait()
|
||||
return _ok_token("bearer-minted-before-invalidation")
|
||||
|
||||
async def re_mint():
|
||||
return _ok_token("bearer-minted-after-invalidation")
|
||||
|
||||
async def must_not_run():
|
||||
pytest.fail("the mint that followed the invalidation should have been cached")
|
||||
|
||||
in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp"))
|
||||
await mint_started.wait()
|
||||
queued = asyncio.create_task(cache.get_or_compute("slot", re_mint, fingerprint="fp"))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert not queued.done()
|
||||
cache.invalidate("slot")
|
||||
release_mint.set()
|
||||
|
||||
raced, fresh = await asyncio.gather(in_flight, queued)
|
||||
assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation"
|
||||
assert isinstance(fresh, Ok) and fresh.ok == "bearer-minted-after-invalidation"
|
||||
|
||||
served = await cache.get_or_compute("slot", must_not_run, fingerprint="fp")
|
||||
assert isinstance(served, Ok) and served.ok == "bearer-minted-after-invalidation"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_does_not_store_a_failed_compute():
|
||||
cache = ExchangedTokenCache()
|
||||
|
|
|
|||
|
|
@ -8198,6 +8198,79 @@ class TestPreemptive401ModeAware:
|
|||
await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False)
|
||||
|
||||
|
||||
def _make_obo_server(alias: str) -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id=f"id-{alias}",
|
||||
name=alias,
|
||||
alias=alias,
|
||||
server_name=alias,
|
||||
url=f"https://{alias}.test/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2_token_exchange,
|
||||
token_exchange_endpoint="https://idp.test/token",
|
||||
client_id="cid",
|
||||
client_secret="csecret",
|
||||
mcp_info={"server_name": alias},
|
||||
)
|
||||
|
||||
|
||||
class TestOboPreflightScopedToAllowedServers:
|
||||
"""The connect-time OBO exchange is an outbound IdP call whose result is cached, so it must
|
||||
only run for a server the caller's key resolves to through the allowed set, not for any
|
||||
server the requested path happens to name."""
|
||||
|
||||
SUBJECT_HEADERS = {"Authorization": "Bearer upstream-subject-token"}
|
||||
|
||||
async def _run(self, requested: MCPServer, allowed: list[MCPServer], user_api_key_auth: UserAPIKeyAuth | None):
|
||||
from litellm.proxy._experimental.mcp_server import server as server_module
|
||||
|
||||
allowed_lookup = AsyncMock(return_value=allowed)
|
||||
preflight = AsyncMock()
|
||||
with (
|
||||
patch.object( # test-quality-ok: route handler reads the module-level manager, no injection seam
|
||||
server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=requested
|
||||
),
|
||||
patch.object( # test-quality-ok: the exchanger is the observable; a real one would call an IdP
|
||||
server_module.global_mcp_server_manager, "preflight_token_exchange", preflight
|
||||
),
|
||||
patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer
|
||||
server_module, "_get_allowed_mcp_servers", allowed_lookup
|
||||
),
|
||||
):
|
||||
await server_module._raise_preemptive_401_for_unauthenticated_servers(
|
||||
scope={"type": "http", "method": "POST", "path": f"/mcp/{requested.alias}", "headers": []},
|
||||
mcp_servers=[requested.alias],
|
||||
oauth2_headers=self.SUBJECT_HEADERS,
|
||||
mcp_server_auth_headers=None,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
client_ip="10.0.0.7",
|
||||
)
|
||||
return allowed_lookup, preflight
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unentitled_key_never_reaches_the_exchanger(self):
|
||||
requested = _make_obo_server("obo_tools")
|
||||
key = UserAPIKeyAuth(api_key="sk-plain-only")
|
||||
|
||||
allowed_lookup, preflight = await self._run(
|
||||
requested, allowed=[_make_obo_server("plain_tools")], user_api_key_auth=key
|
||||
)
|
||||
|
||||
preflight.assert_not_awaited()
|
||||
allowed_lookup.assert_awaited_once_with(
|
||||
user_api_key_auth=key, mcp_servers=[requested.alias], client_ip="10.0.0.7"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entitled_key_still_exchanges_at_connect(self):
|
||||
requested = _make_obo_server("obo_tools")
|
||||
key = UserAPIKeyAuth(api_key="sk-obo")
|
||||
|
||||
_, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key)
|
||||
|
||||
preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_guardrails_return_the_rewritten_result():
|
||||
"""The result a post_mcp_call guardrail rewrote must be what the caller sends back."""
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ def _mock_cli_sso_start_response(
|
|||
login_id: str = "cli-session-uuid-456",
|
||||
poll_secret: str = "poll-secret",
|
||||
user_code: str = "ABCD-EFGH",
|
||||
**extra_fields: object,
|
||||
) -> Mock:
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
|
|
@ -66,6 +67,7 @@ def _mock_cli_sso_start_response(
|
|||
"login_id": login_id,
|
||||
"poll_secret": poll_secret,
|
||||
"user_code": user_code,
|
||||
**extra_fields,
|
||||
}
|
||||
mock_response.raise_for_status = Mock()
|
||||
return mock_response
|
||||
|
|
@ -333,7 +335,9 @@ class TestLoginCommand:
|
|||
call_args = mock_browser.call_args[0][0]
|
||||
assert "https://test.example.com/sso/key/generate" in call_args
|
||||
assert "cli-test-uuid-123" in call_args
|
||||
assert "user_code" not in call_args
|
||||
assert "Verification code: ABCD-EFGH" in result.output
|
||||
assert "pre-filled in the browser" not in result.output
|
||||
mock_post.assert_called_once()
|
||||
mock_get.assert_called()
|
||||
assert mock_get.call_args.kwargs["headers"] == {"x-litellm-cli-poll-secret": "poll-secret"}
|
||||
|
|
@ -347,6 +351,72 @@ class TestLoginCommand:
|
|||
# Verify commands were shown
|
||||
mock_show_commands.assert_called_once()
|
||||
|
||||
def test_login_prefills_the_code_in_the_browser_when_the_proxy_advertises_it(
|
||||
self, isolated_home, secret_vault_factory
|
||||
) -> None:
|
||||
vault = secret_vault_factory()
|
||||
poll_response = Mock()
|
||||
poll_response.status_code = 200
|
||||
poll_response.json.return_value = {
|
||||
"status": "ready",
|
||||
"key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt",
|
||||
"user_id": "test-user-123",
|
||||
"team_id": "team-1",
|
||||
"teams": ["team-1"],
|
||||
}
|
||||
start_response = _mock_cli_sso_start_response(
|
||||
login_id="cli-test-uuid-123",
|
||||
verification_uri_complete=(
|
||||
"https://internal-hostname.example.com/sso/key/generate"
|
||||
"?source=litellm-cli&key=cli-test-uuid-123&user_code=ABCD-EFGH"
|
||||
),
|
||||
)
|
||||
|
||||
with (
|
||||
patch("webbrowser.open") as mock_browser,
|
||||
patch("requests.post", return_value=start_response),
|
||||
patch("requests.get", return_value=poll_response),
|
||||
):
|
||||
result = self.runner.invoke(login, obj={"base_url": "https://test.example.com", "secret_vault": vault})
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(vault.blob)["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt"
|
||||
assert json.loads((isolated_home / ".litellm" / "token.json").read_text())["user_id"] == "test-user-123"
|
||||
opened_url = mock_browser.call_args[0][0]
|
||||
assert opened_url.startswith("https://test.example.com/sso/key/generate?")
|
||||
assert "internal-hostname" not in opened_url
|
||||
assert "key=cli-test-uuid-123" in opened_url
|
||||
assert "user_code=ABCD-EFGH" in opened_url
|
||||
assert "Verification code: ABCD-EFGH (pre-filled in the browser, check it matches)" in result.output
|
||||
|
||||
def test_login_keeps_the_code_out_of_the_url_when_the_proxy_sends_a_non_url_verification_uri(
|
||||
self, secret_vault_factory
|
||||
) -> None:
|
||||
poll_response = Mock()
|
||||
poll_response.status_code = 200
|
||||
poll_response.json.return_value = {
|
||||
"status": "ready",
|
||||
"key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt",
|
||||
"user_id": "test-user-123",
|
||||
"team_id": "team-1",
|
||||
"teams": ["team-1"],
|
||||
}
|
||||
for advertised in (None, True):
|
||||
start_response = _mock_cli_sso_start_response(verification_uri_complete=advertised)
|
||||
|
||||
with (
|
||||
patch("webbrowser.open") as mock_browser,
|
||||
patch("requests.post", return_value=start_response),
|
||||
patch("requests.get", return_value=poll_response),
|
||||
):
|
||||
result = self.runner.invoke(
|
||||
login, obj={"base_url": "https://test.example.com", "secret_vault": secret_vault_factory()}
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "user_code" not in mock_browser.call_args[0][0]
|
||||
assert "pre-filled in the browser" not in result.output
|
||||
|
||||
def test_login_timeout(self):
|
||||
"""Test login timeout scenario"""
|
||||
mock_context = Mock()
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ from unittest.mock import patch
|
|||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import Timeout
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import initialize_guardrail
|
||||
|
|
@ -12,8 +14,8 @@ from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr
|
|||
CrowdStrikeAIDRHandler,
|
||||
)
|
||||
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse
|
||||
from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams
|
||||
from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -1578,3 +1580,139 @@ async def test_unparseable_transformed_response_fails_closed_under_fail_open() -
|
|||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "failing closed" in exc_info.value.detail["error"]
|
||||
|
||||
|
||||
def _initialize_from_config(**litellm_params_kwargs: object) -> CrowdStrikeAIDRHandler:
|
||||
litellm_params = LitellmParams(
|
||||
guardrail="crowdstrike_aidr",
|
||||
api_key="pts_crowdstrike_tokenid",
|
||||
api_base="https://api.crowdstrike.com/aidr/aiguard",
|
||||
default_on=True,
|
||||
**litellm_params_kwargs,
|
||||
)
|
||||
guardrail = Guardrail(guardrail_name="crowdstrike-aidr-guard", litellm_params=litellm_params)
|
||||
return initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "runs_pre_call", "runs_post_call"),
|
||||
[("post_call", False, True), ("pre_call", True, False), (["pre_call", "post_call"], True, True)],
|
||||
)
|
||||
def test_initialize_guardrail_honors_configured_mode(
|
||||
mode: str | list[str], runs_pre_call: bool, runs_post_call: bool
|
||||
) -> None:
|
||||
handler = _initialize_from_config(mode=mode)
|
||||
|
||||
assert handler.should_run_guardrail({}, GuardrailEventHooks.pre_call) is runs_pre_call
|
||||
assert handler.should_run_guardrail({}, GuardrailEventHooks.post_call) is runs_post_call
|
||||
|
||||
|
||||
def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_hooks() -> None:
|
||||
with pytest.raises(ValueError, match="during_call is not in the supported event hooks"):
|
||||
_initialize_from_config(mode="during_call")
|
||||
|
||||
|
||||
def test_initialize_guardrail_defaults_streaming_params() -> None:
|
||||
handler = _initialize_from_config(mode="post_call")
|
||||
|
||||
assert handler.streaming_end_of_stream_only is False
|
||||
assert handler.streaming_sampling_rate == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"configured",
|
||||
[
|
||||
{"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50},
|
||||
{"optional_params": {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50}},
|
||||
],
|
||||
)
|
||||
def test_initialize_guardrail_forwards_streaming_params(configured: dict[str, object]) -> None:
|
||||
handler = _initialize_from_config(mode="post_call", **configured)
|
||||
|
||||
assert handler.streaming_end_of_stream_only is True
|
||||
assert handler.streaming_sampling_rate == 50
|
||||
|
||||
|
||||
def test_initialize_guardrail_rejects_non_positive_sampling_rate() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
_initialize_from_config(mode="post_call", streaming_sampling_rate=0)
|
||||
|
||||
|
||||
def test_update_in_memory_litellm_params_reapplies_streaming_params() -> None:
|
||||
handler = _initialize_from_config(mode="post_call")
|
||||
|
||||
handler.update_in_memory_litellm_params(
|
||||
LitellmParams(
|
||||
guardrail="crowdstrike_aidr",
|
||||
mode="post_call",
|
||||
streaming_end_of_stream_only=True,
|
||||
streaming_sampling_rate=7,
|
||||
)
|
||||
)
|
||||
|
||||
assert handler.streaming_end_of_stream_only is True
|
||||
assert handler.streaming_sampling_rate == 7
|
||||
|
||||
|
||||
def _stream_chunk(content: str, finish_reason: str | None) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
model="gpt-4",
|
||||
choices=[
|
||||
litellm.StreamingChoices(
|
||||
index=0, delta=Delta(role="assistant", content=content), finish_reason=finish_reason
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts: list[str]) -> int:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
|
||||
|
||||
async def stream():
|
||||
for i, content in enumerate(chunk_texts):
|
||||
yield _stream_chunk(content, "stop" if i == len(chunk_texts) - 1 else None)
|
||||
|
||||
calls = 0
|
||||
|
||||
def _allow(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return httpx.Response(
|
||||
status_code=200, json={"result": {"blocked": False, "transformed": False}}, request=request
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": handler,
|
||||
"metadata": {"guardrails": ["crowdstrike-aidr-guard"]},
|
||||
}
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(_allow)) as client:
|
||||
await handler.async_handler.close()
|
||||
handler.async_handler.client = client
|
||||
async for _ in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"),
|
||||
response=stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
return calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("configured", "expected_calls"),
|
||||
[
|
||||
({}, 3),
|
||||
({"streaming_sampling_rate": 2}, 6),
|
||||
({"streaming_end_of_stream_only": True}, 1),
|
||||
({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1),
|
||||
],
|
||||
)
|
||||
async def test_streaming_params_from_config_control_output_scan_cadence(
|
||||
configured: dict[str, object], expected_calls: int
|
||||
) -> None:
|
||||
"""10 chunks: default samples at 5 and 10 plus the final pass, rate 2 samples 5 times plus final, end-of-stream scans once."""
|
||||
handler = _initialize_from_config(mode="post_call", **configured)
|
||||
|
||||
assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls
|
||||
|
|
|
|||
|
|
@ -2854,10 +2854,12 @@ async def test_ui_view_request_response_custom_logger_is_keyed_by_callers_own_re
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_id_lookup_lists_exact_request_id_row_first(client, monkeypatch):
|
||||
@pytest.mark.parametrize("group_by_session", [False, True])
|
||||
async def test_ui_view_spend_logs_id_lookup_lists_exact_request_id_row_first(client, monkeypatch, group_by_session):
|
||||
"""The dashboard's deep link fetches a single row for ``?log_id=``. When a newer
|
||||
row carries that id as its client-set litellm_call_id, the row whose request_id
|
||||
is the id still comes first, so the link opens the request it names."""
|
||||
is the id still comes first, so the link opens the request it names. The
|
||||
session-grouped page orders its representatives the same way."""
|
||||
today = datetime.datetime.now(timezone.utc)
|
||||
corpus = [
|
||||
{
|
||||
|
|
@ -2898,7 +2900,7 @@ async def test_ui_view_spend_logs_id_lookup_lists_exact_request_id_row_first(cli
|
|||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"request_id": "victim-req", "page_size": 1},
|
||||
params={"request_id": "victim-req", "page_size": 1, "group_by_session": str(group_by_session).lower()},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
|
@ -4580,18 +4582,18 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
|
|||
]
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_spendlogs.group_by = AsyncMock(
|
||||
return_value=[
|
||||
{"session_id": session_id, "_count": {"session_id": 2}},
|
||||
]
|
||||
)
|
||||
mock_prisma.db.litellm_spendlogs.group_by = AsyncMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"session_id": session_id,
|
||||
"api_key": api_key,
|
||||
"session_total_count": 2,
|
||||
"session_total_spend": 15.0,
|
||||
"mcp_tool_call_count": 1,
|
||||
"mcp_tool_call_spend": 10.0,
|
||||
"session_llm_count": 1,
|
||||
"session_agent_count": 0,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
|
@ -4616,6 +4618,8 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
|
|||
assert rows[0]["mcp_tool_call_spend"] == 10.0
|
||||
assert rows[1]["mcp_tool_call_count"] == 1
|
||||
assert rows[1]["mcp_tool_call_spend"] == 10.0
|
||||
assert rows[0]["session_llm_count"] == 1
|
||||
assert rows[0]["session_agent_count"] == 0
|
||||
|
||||
# Every row in the session carries the full session spend, not just its own
|
||||
assert rows[0]["session_total_spend"] == 15.0
|
||||
|
|
@ -4624,13 +4628,126 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
|
|||
# Row without a session_id defaults to 1
|
||||
assert rows[2]["session_total_count"] == 1
|
||||
|
||||
# group_by should have been called with the session_id
|
||||
mock_prisma.db.litellm_spendlogs.group_by.assert_called_once_with(
|
||||
by=["session_id"],
|
||||
where={"session_id": {"in": [session_id]}},
|
||||
count={"session_id": True},
|
||||
# The count is folded into the single aggregate query; no separate group_by call.
|
||||
mock_prisma.db.litellm_spendlogs.group_by.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_ui_spend_logs_response_key_split_session_gets_per_key_aggregates():
|
||||
"""
|
||||
Two keys reusing one session id are separate rows under grouped pagination,
|
||||
and each row must carry ITS key's totals, never the combined session's:
|
||||
the aggregate query and its lookup are keyed by (session_id, api_key).
|
||||
"""
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
_build_ui_spend_logs_response,
|
||||
)
|
||||
|
||||
session_id = "sess-shared"
|
||||
dict_rows = [
|
||||
{"request_id": "req-a", "session_id": session_id, "call_type": "completion", "api_key": "key-a"},
|
||||
{"request_id": "req-b", "session_id": session_id, "call_type": "completion", "api_key": "key-b"},
|
||||
]
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"session_id": session_id,
|
||||
"api_key": "key-a",
|
||||
"session_total_count": 2,
|
||||
"session_total_spend": 0.2,
|
||||
"mcp_tool_call_count": 0,
|
||||
"mcp_tool_call_spend": 0.0,
|
||||
"session_cache_hit_count": 1,
|
||||
"session_llm_count": 2,
|
||||
"session_agent_count": 0,
|
||||
},
|
||||
{
|
||||
"session_id": session_id,
|
||||
"api_key": "key-b",
|
||||
"session_total_count": 1,
|
||||
"session_total_spend": 0.7,
|
||||
"mcp_tool_call_count": 0,
|
||||
"mcp_tool_call_spend": 0.0,
|
||||
"session_cache_hit_count": 0,
|
||||
"session_llm_count": 1,
|
||||
"session_agent_count": 0,
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
result = await _build_ui_spend_logs_response(
|
||||
prisma_client=mock_prisma,
|
||||
data=dict_rows,
|
||||
total_records=2,
|
||||
page=1,
|
||||
page_size=50,
|
||||
total_pages=1,
|
||||
enrich_session_counts=True,
|
||||
)
|
||||
|
||||
rows = result["data"]
|
||||
assert [(r["session_total_count"], r["session_total_spend"]) for r in rows] == [(2, 0.2), (1, 0.7)]
|
||||
assert [r["session_cache_hit_count"] for r in rows] == [1, 0]
|
||||
assert [r["session_llm_count"] for r in rows] == [2, 1]
|
||||
|
||||
aggregate_sql = mock_prisma.db.query_raw.mock_calls[0][1][0]
|
||||
assert "GROUP BY session_id, api_key" in aggregate_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_ui_spend_logs_response_empty_api_key_keeps_session_aggregates():
|
||||
"""
|
||||
The spend-log schema defaults api_key to an empty string, which is a real
|
||||
group value and not a missing one: a multi-call session logged under an
|
||||
empty key must keep its count and spend instead of degrading to a plain
|
||||
single-call row.
|
||||
"""
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
_build_ui_spend_logs_response,
|
||||
)
|
||||
|
||||
session_id = "sess-keyless"
|
||||
dict_rows = [
|
||||
{"request_id": "req-1", "session_id": session_id, "call_type": "completion", "api_key": ""},
|
||||
]
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"session_id": session_id,
|
||||
"api_key": "",
|
||||
"session_total_count": 3,
|
||||
"session_total_spend": 0.09,
|
||||
"mcp_tool_call_count": 0,
|
||||
"mcp_tool_call_spend": 0.0,
|
||||
"session_cache_hit_count": 0,
|
||||
"session_llm_count": 3,
|
||||
"session_agent_count": 0,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
result = await _build_ui_spend_logs_response(
|
||||
prisma_client=mock_prisma,
|
||||
data=dict_rows,
|
||||
total_records=1,
|
||||
page=1,
|
||||
page_size=50,
|
||||
total_pages=1,
|
||||
enrich_session_counts=True,
|
||||
)
|
||||
|
||||
row = result["data"][0]
|
||||
assert row["session_total_count"] == 3
|
||||
assert row["session_total_spend"] == 0.09
|
||||
|
||||
# The empty key must reach the aggregate's authorized-keys filter too.
|
||||
_, call_args, _ = mock_prisma.db.query_raw.mock_calls[0]
|
||||
assert call_args[2] == [""]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_ui_spend_logs_response_sums_multi_round_session_spend():
|
||||
|
|
@ -4654,14 +4771,13 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend():
|
|||
]
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_spendlogs.group_by = AsyncMock(
|
||||
return_value=[{"session_id": session_id, "_count": {"session_id": 3}}]
|
||||
)
|
||||
# The raw aggregate query returns the full session spend (0.01 + 0.02 + 0.03).
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"session_id": session_id,
|
||||
"api_key": api_key,
|
||||
"session_total_count": 3,
|
||||
"session_total_spend": 0.06,
|
||||
"mcp_tool_call_count": 0,
|
||||
"mcp_tool_call_spend": 0.0,
|
||||
|
|
@ -4710,13 +4826,12 @@ async def test_build_ui_spend_logs_response_session_cache_hit_count():
|
|||
]
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_spendlogs.group_by = AsyncMock(
|
||||
return_value=[{"session_id": session_id, "_count": {"session_id": 2}}]
|
||||
)
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"session_id": session_id,
|
||||
"api_key": api_key,
|
||||
"session_total_count": 2,
|
||||
"session_total_spend": 0.05,
|
||||
"mcp_tool_call_count": 0,
|
||||
"mcp_tool_call_spend": 0.0,
|
||||
|
|
|
|||
|
|
@ -274,6 +274,9 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch):
|
|||
"the page query must not carry a window count that forces a full-window "
|
||||
f"scan. SQL was:\n{page_sql}"
|
||||
)
|
||||
assert "GROUP BY" not in count_sql and "DISTINCT ON" not in page_sql, (
|
||||
"without group_by_session the endpoint must keep raw per-call pagination"
|
||||
)
|
||||
|
||||
assert response["total"] == 137
|
||||
assert response["total_is_capped"] is False
|
||||
|
|
@ -499,3 +502,106 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch):
|
|||
params = mock_prisma.db.query_raw.call_args[0][1:]
|
||||
assert "team_x" in params, "team_id must be forwarded into the DB query params"
|
||||
assert "sl.team_id = $3" in sql, f"team query must filter on team_id. SQL was:\n{sql}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
|
||||
"""
|
||||
With group_by_session=true, /spend/logs/ui must page and count SESSIONS,
|
||||
not raw calls: the page query returns one representative row per session
|
||||
(DISTINCT ON the session group key, preferring non-MCP calls, newest
|
||||
first) and the bounded count counts groups. Otherwise the UI collapses a
|
||||
server page of N calls into fewer visible rows while the footer still
|
||||
claims N (issue #38060).
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
SPEND_LOGS_PAGINATION_COUNT_CAP,
|
||||
ui_view_spend_logs,
|
||||
)
|
||||
|
||||
page_rows = [
|
||||
{"request_id": "req-1", "metadata": "{}", "session_id": None},
|
||||
{"request_id": "req-2", "metadata": "{}", "session_id": None},
|
||||
]
|
||||
mock_prisma = _make_ui_spend_logs_mock(count_total=12, page_rows=page_rows)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/spend/logs/ui"
|
||||
|
||||
response = await ui_view_spend_logs(
|
||||
request=mock_request,
|
||||
api_key=None,
|
||||
user_id=None,
|
||||
request_id=None,
|
||||
start_date="2026-02-16 00:00:00",
|
||||
end_date="2026-02-16 23:59:59",
|
||||
page=1,
|
||||
page_size=50,
|
||||
sort_by="startTime",
|
||||
sort_order="desc",
|
||||
user_api_key_dict=auth,
|
||||
group_by_session=True,
|
||||
)
|
||||
|
||||
group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
|
||||
|
||||
count_call = mock_prisma.db.query_raw.call_args_list[0]
|
||||
count_sql = count_call[0][0]
|
||||
assert f"GROUP BY {group_key}" in count_sql, f"grouped total must count sessions. SQL was:\n{count_sql}"
|
||||
assert "COUNT(*) OVER ()" not in count_sql
|
||||
assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded"
|
||||
assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1
|
||||
|
||||
page_sql = mock_prisma.db.query_raw.call_args_list[1][0][0]
|
||||
assert f"DISTINCT ON ({group_key})" in page_sql, f"page must return one row per session. SQL was:\n{page_sql}"
|
||||
assert f"ORDER BY {group_key}, call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC" in page_sql, (
|
||||
"the session representative must prefer the newest non-MCP call"
|
||||
)
|
||||
assert "COUNT(*) OVER ()" not in page_sql
|
||||
|
||||
assert response["total"] == 12
|
||||
assert response["total_is_capped"] is False
|
||||
assert response["total_pages"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(monkeypatch):
|
||||
"""
|
||||
A request_id lookup with group_by_session=true must still resolve the
|
||||
exact requested row: the filter runs before grouping, so the row is its
|
||||
own group's representative and deep links keep working.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs
|
||||
|
||||
target_row = {"request_id": "req-deep-link", "metadata": "{}", "session_id": None}
|
||||
mock_prisma = _make_ui_spend_logs_mock(count_total=1, page_rows=[target_row])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/spend/logs/ui"
|
||||
|
||||
response = await ui_view_spend_logs(
|
||||
request=mock_request,
|
||||
api_key=None,
|
||||
user_id=None,
|
||||
request_id="req-deep-link",
|
||||
start_date=None,
|
||||
end_date=None,
|
||||
page=1,
|
||||
page_size=1,
|
||||
sort_by="startTime",
|
||||
sort_order="desc",
|
||||
user_api_key_dict=auth,
|
||||
group_by_session=True,
|
||||
)
|
||||
|
||||
page_call = mock_prisma.db.query_raw.call_args_list[1]
|
||||
assert "request_id = $" in page_call[0][0], "the request_id equality filter must survive grouping"
|
||||
assert "req-deep-link" in page_call[0]
|
||||
assert [row["request_id"] for row in response["data"]] == ["req-deep-link"]
|
||||
assert response["total"] == 1
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ completion_start_time = end_time."""
|
|||
import json
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
from unittest.mock import Mock
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -378,3 +378,162 @@ def test_stamp_responses_usage_cost_survives_calculator_failure():
|
|||
_stamp_responses_usage_cost(response, logging_obj)
|
||||
|
||||
assert getattr(response.usage, "cost", None) is None
|
||||
|
||||
|
||||
def _capture_dispatch(logged: list):
|
||||
"""Record the object handed to the success handlers.
|
||||
|
||||
``Mock(spec=LiteLLMLoggingObj).dispatch_success_handlers`` is an AsyncMock whose side effect
|
||||
only runs when the coroutine is awaited, so capture with a plain function instead.
|
||||
"""
|
||||
|
||||
async def _noop() -> None:
|
||||
return None
|
||||
|
||||
def _dispatch(result, **kwargs):
|
||||
logged.append(result)
|
||||
return _noop()
|
||||
|
||||
return _dispatch
|
||||
|
||||
|
||||
def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock:
|
||||
"""Config whose completed event carries a real ResponsesAPIResponse, so the logging copy
|
||||
performs a genuine model_dump/model_validate round trip."""
|
||||
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
||||
|
||||
def _transform(model, parsed_chunk, logging_obj):
|
||||
evt_type = parsed_chunk.get("type")
|
||||
if evt_type != "response.completed":
|
||||
stub = Mock()
|
||||
stub.type = evt_type
|
||||
return stub
|
||||
response = ResponsesAPIResponse(
|
||||
id="resp_headers",
|
||||
created_at=1,
|
||||
output=[],
|
||||
parallel_tool_calls=False,
|
||||
tool_choice="auto",
|
||||
tools=[],
|
||||
)
|
||||
if transform_hidden_params is not None:
|
||||
response._hidden_params.update(transform_hidden_params)
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=response,
|
||||
)
|
||||
|
||||
mock_config.transform_streaming_response.side_effect = _transform
|
||||
return mock_config
|
||||
|
||||
|
||||
def _make_header_iterator(
|
||||
*,
|
||||
headers: dict,
|
||||
config: Mock,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> ResponsesAPIStreamingIterator:
|
||||
async def aiter_bytes():
|
||||
yield _sse_event({"type": "response.completed"})
|
||||
|
||||
mock_response = Mock()
|
||||
mock_response.headers = headers
|
||||
mock_response.aiter_bytes = aiter_bytes
|
||||
|
||||
return ResponsesAPIStreamingIterator(
|
||||
response=mock_response,
|
||||
model="gpt-4o-mini",
|
||||
responses_api_provider_config=config,
|
||||
logging_obj=logging_obj,
|
||||
litellm_metadata={},
|
||||
custom_llm_provider="azure",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_logging_response_carries_provider_response_headers():
|
||||
"""LIT-6055: the provider headers the iterator captured must reach the logged response, so
|
||||
custom loggers can read Azure's apim-request-id from the callback payload."""
|
||||
logging_obj = _logging_obj_stub()
|
||||
logged: list[object] = []
|
||||
logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
|
||||
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
|
||||
iterator = _make_header_iterator(
|
||||
headers={"apim-request-id": "azure-correlation-1", "x-ms-region": "East US 2"},
|
||||
config=_headers_config(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
assert len(logged) == 1
|
||||
hidden_params = logged[0].response._hidden_params
|
||||
assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "azure-correlation-1"
|
||||
assert hidden_params["additional_headers"]["llm_provider-x-ms-region"] == "East US 2"
|
||||
assert hidden_params["headers"]["apim-request-id"] == "azure-correlation-1"
|
||||
# the proxy builds the client's response headers from the iterator's own dict, so the logged
|
||||
# response must hold copies rather than alias it
|
||||
assert hidden_params["additional_headers"] is not iterator._hidden_params["additional_headers"]
|
||||
assert hidden_params["headers"] is not iterator._raw_response_headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_logging_copy_preserves_transform_hidden_params():
|
||||
"""LIT-6055: model_validate(model_dump()) drops pydantic private attributes, so headers a
|
||||
provider transform already set on the response (fake_stream) must be re-applied."""
|
||||
logging_obj = _logging_obj_stub()
|
||||
logged: list[object] = []
|
||||
logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
|
||||
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
|
||||
iterator = _make_header_iterator(
|
||||
headers={},
|
||||
config=_headers_config(
|
||||
transform_hidden_params={
|
||||
"additional_headers": {"llm_provider-apim-request-id": "from-transform"},
|
||||
"headers": {"apim-request-id": "from-transform"},
|
||||
"response_cost": 0.5,
|
||||
}
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
assert len(logged) == 1
|
||||
hidden_params = logged[0].response._hidden_params
|
||||
assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "from-transform"
|
||||
assert hidden_params["headers"]["apim-request-id"] == "from-transform"
|
||||
assert iterator.completed_response is not logged[0]
|
||||
# only the header keys travel: response_cost would short-circuit the cost calculator
|
||||
assert "response_cost" not in hidden_params
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched():
|
||||
"""LIT-6055: when the logging copy falls back to the original event, the header restore must
|
||||
not stamp logging-only state onto the object the caller is iterating."""
|
||||
logging_obj = _logging_obj_stub()
|
||||
logged: list[object] = []
|
||||
logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
|
||||
iterator = _make_header_iterator(
|
||||
headers={"apim-request-id": "azure-correlation-1"},
|
||||
config=_headers_config(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
async for _ in iterator:
|
||||
pass
|
||||
|
||||
assert len(logged) == 1
|
||||
iterator._completed_response_logged = False
|
||||
logged.clear()
|
||||
with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")):
|
||||
iterator._log_completed_response(is_async=True)
|
||||
|
||||
assert logged == [iterator.completed_response]
|
||||
assert iterator.completed_response.response._hidden_params == {}
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 22334
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26765
|
||||
"limit": 26763
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 261
|
||||
|
|
@ -27,12 +27,12 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16482
|
||||
"limit": 16480
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5520
|
||||
},
|
||||
"LIT012": {
|
||||
"limit": 4495
|
||||
"limit": 4489
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2042,6 +2042,7 @@ interface UiSpendLogsParams {
|
|||
min_spend?: number;
|
||||
max_spend?: number;
|
||||
exclude_internal_health_checks?: boolean;
|
||||
group_by_session?: boolean;
|
||||
}
|
||||
|
||||
interface UiSpendLogsCallOptions {
|
||||
|
|
|
|||
|
|
@ -138,33 +138,39 @@ describe("RequestLogsPanel", () => {
|
|||
respondWith([]);
|
||||
});
|
||||
|
||||
describe("multi-call session collapsing", () => {
|
||||
const sessionRows = [
|
||||
logEntry({ request_id: "req-mcp", call_type: "call_mcp_tool", session_id: "sess-1", session_total_count: 3 }),
|
||||
logEntry({ request_id: "req-llm", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }),
|
||||
logEntry({ request_id: "req-llm-2", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }),
|
||||
];
|
||||
|
||||
it("collapses a multi-call session to a single representative row", async () => {
|
||||
respondWith(sessionRows);
|
||||
describe("server-grouped session pagination (#38060)", () => {
|
||||
it("requests session-grouped pages of 10 rows by default", async () => {
|
||||
renderPanel();
|
||||
|
||||
await waitFor(() => expect(row("req-mcp") ?? row("req-llm") ?? row("req-llm-2")).not.toBeNull());
|
||||
|
||||
const rendered = ["req-mcp", "req-llm", "req-llm-2"].filter((id) => row(id) !== null);
|
||||
expect(rendered).toHaveLength(1);
|
||||
await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
|
||||
expect(lastCall()?.params?.group_by_session).toBe(true);
|
||||
expect(lastCall()?.page_size).toBe(10);
|
||||
});
|
||||
|
||||
it("prefers an LLM call over an MCP call as the session's representative", async () => {
|
||||
respondWith(sessionRows);
|
||||
it("renders every row the server returns without client-side collapsing", async () => {
|
||||
respondWith([
|
||||
logEntry({ request_id: "req-a", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }),
|
||||
logEntry({ request_id: "req-b", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }),
|
||||
logEntry({ request_id: "req-c", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }),
|
||||
]);
|
||||
renderPanel();
|
||||
|
||||
await waitFor(() => expect(row("req-llm")).not.toBeNull());
|
||||
expect(row("req-mcp")).toBeNull();
|
||||
await waitFor(() => expect(row("req-a")).not.toBeNull());
|
||||
expect(row("req-b")).not.toBeNull();
|
||||
expect(row("req-c")).not.toBeNull();
|
||||
});
|
||||
|
||||
it("shows the session's call count and composition on the representative row", async () => {
|
||||
respondWith(sessionRows);
|
||||
it("shows the session's call count on the server-picked representative row", async () => {
|
||||
respondWith([
|
||||
logEntry({
|
||||
request_id: "req-llm",
|
||||
call_type: "acompletion",
|
||||
session_id: "sess-1",
|
||||
session_total_count: 3,
|
||||
session_llm_count: 2,
|
||||
mcp_tool_call_count: 1,
|
||||
}),
|
||||
]);
|
||||
renderPanel();
|
||||
|
||||
await waitFor(() => expect(row("req-llm")).not.toBeNull());
|
||||
|
|
@ -296,6 +302,7 @@ describe("RequestLogsPanel", () => {
|
|||
if (!byIdCall) throw new Error("expected a by-id uiSpendLogsCall");
|
||||
expect(byIdCall.page).toBe(1);
|
||||
expect(byIdCall.page_size).toBe(1);
|
||||
expect(byIdCall.params?.group_by_session).toBeUndefined();
|
||||
});
|
||||
|
||||
it("opens the drawer when ?log_id= is the log's litellm_call_id rather than its request_id", async () => {
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import type { KeyResponse } from "../key_team_helpers/key_list";
|
|||
import { keyInfoV1Call, uiSpendLogsCall } from "../networking";
|
||||
import KeyInfoView from "../templates/key_info_view";
|
||||
import type { LogEntry } from "./columns";
|
||||
import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "./constants";
|
||||
import { LOGS_PAGE_SIZE_OPTIONS } from "./constants";
|
||||
import {
|
||||
DEFAULT_LOGS_SORTING,
|
||||
formatLogsWindow,
|
||||
|
|
@ -24,7 +24,7 @@ import { LogDetailsDrawer } from "./LogDetailsDrawer";
|
|||
import { LiveTailBanner, LogsTableToolbar } from "./LogsTableToolbar";
|
||||
import { RequestLogsTable } from "./RequestLogsTable";
|
||||
|
||||
const PAGE_SIZE = 50;
|
||||
const PAGE_SIZE = LOGS_PAGE_SIZE_OPTIONS[0];
|
||||
const DEFAULT_INTERVAL = { value: 24, unit: "hours" };
|
||||
const matchesLogId = (log: LogEntry, logId: string) => log.request_id === logId || log.litellm_call_id === logId;
|
||||
const findLogById = (logs: readonly LogEntry[], logId: string): LogEntry | null =>
|
||||
|
|
@ -38,12 +38,6 @@ interface RequestLogsPanelProps {
|
|||
isActive: boolean;
|
||||
}
|
||||
|
||||
interface SessionComposition {
|
||||
llm: number;
|
||||
agent: number;
|
||||
mcp: number;
|
||||
}
|
||||
|
||||
export default function RequestLogsPanel({ accessToken, token, userRole, userID, isActive }: RequestLogsPanelProps) {
|
||||
const [pagination, setPagination] = useState<PaginationState>({ pageIndex: 0, pageSize: PAGE_SIZE });
|
||||
const [sorting, setSorting] = useState<SortingState>(DEFAULT_LOGS_SORTING);
|
||||
|
|
@ -160,49 +154,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
|
||||
const isDrawerOpen = displayLog !== null || displaySessionId !== null;
|
||||
|
||||
const rows = useMemo<LogEntry[]>(() => {
|
||||
const searchedLogs = filteredLogs.data;
|
||||
|
||||
const sessionCompositionById = searchedLogs.reduce<Record<string, SessionComposition>>((acc, log) => {
|
||||
if (!log.session_id) return acc;
|
||||
if (!acc[log.session_id]) {
|
||||
acc[log.session_id] = { llm: 0, agent: 0, mcp: 0 };
|
||||
}
|
||||
if (MCP_CALL_TYPES.includes(log.call_type)) {
|
||||
acc[log.session_id].mcp += 1;
|
||||
} else if (AGENT_CALL_TYPES.includes(log.call_type)) {
|
||||
acc[log.session_id].agent += 1;
|
||||
} else {
|
||||
acc[log.session_id].llm += 1;
|
||||
}
|
||||
return acc;
|
||||
}, {});
|
||||
|
||||
const sessionRepresentativeMap = new Map<string, { requestId: string; isMcp: boolean }>();
|
||||
for (const log of searchedLogs) {
|
||||
if (!log.session_id || (log.session_total_count || 1) <= 1) continue;
|
||||
const isMcp = MCP_CALL_TYPES.includes(log.call_type);
|
||||
const existing = sessionRepresentativeMap.get(log.session_id);
|
||||
if (!existing || (existing.isMcp && !isMcp)) {
|
||||
sessionRepresentativeMap.set(log.session_id, { requestId: log.request_id, isMcp });
|
||||
}
|
||||
}
|
||||
|
||||
return searchedLogs
|
||||
.map((log) => {
|
||||
const sessionComposition = log.session_id ? sessionCompositionById[log.session_id] : undefined;
|
||||
return {
|
||||
...log,
|
||||
session_llm_count: sessionComposition?.llm ?? undefined,
|
||||
session_mcp_count: sessionComposition?.mcp ?? undefined,
|
||||
session_agent_count: sessionComposition?.agent ?? undefined,
|
||||
};
|
||||
})
|
||||
.filter((log) => {
|
||||
if (!log.session_id || (log.session_total_count || 1) <= 1) return true;
|
||||
return sessionRepresentativeMap.get(log.session_id)?.requestId === log.request_id;
|
||||
});
|
||||
}, [filteredLogs.data]);
|
||||
const rows: LogEntry[] = filteredLogs.data;
|
||||
|
||||
const searchTerm = useMemo(() => {
|
||||
const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.REQUEST_ID);
|
||||
|
|
@ -261,13 +213,12 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
);
|
||||
|
||||
const handleSessionClick = useCallback(
|
||||
(sessionId: string) => {
|
||||
if (!sessionId) return;
|
||||
const log = rows.find((candidate) => candidate.session_id === sessionId) ?? null;
|
||||
(log: LogEntry) => {
|
||||
if (!log.session_id) return;
|
||||
setSelectedLog(log);
|
||||
openSession(sessionId, log?.request_id ?? null);
|
||||
openSession(log.session_id, log.request_id);
|
||||
},
|
||||
[rows, openSession],
|
||||
[openSession],
|
||||
);
|
||||
|
||||
const handleSelectLog = useCallback(
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import { DataTable, DataTableFilterDrawer, DataTableToolbar } from "@/components
|
|||
|
||||
import type { Team } from "../key_team_helpers/key_list";
|
||||
import type { LogEntry } from "./columns";
|
||||
import { LOGS_PAGE_SIZE_OPTIONS } from "./constants";
|
||||
import { LOG_FILTER_LABELS, type LogsWindow } from "./log_filter_logic";
|
||||
import { RequestLogsFilters } from "./RequestLogsFilters";
|
||||
import { getRequestLogsTableColumns } from "./RequestLogsTableColumns";
|
||||
|
|
@ -28,7 +29,7 @@ interface RequestLogsTableProps {
|
|||
onRefresh: () => void;
|
||||
onRowClick: (log: LogEntry) => void;
|
||||
onKeyHashClick: (keyHash: string) => void;
|
||||
onSessionClick: (sessionId: string) => void;
|
||||
onSessionClick: (log: LogEntry) => void;
|
||||
teams: Team[];
|
||||
logsWindow: LogsWindow;
|
||||
toolbarChildren?: ReactNode;
|
||||
|
|
@ -91,6 +92,7 @@ export function RequestLogsTable({
|
|||
paginationMode="server"
|
||||
pagination={pagination}
|
||||
onPaginationChange={onPaginationChange}
|
||||
pageSizeOptions={LOGS_PAGE_SIZE_OPTIONS}
|
||||
rowCount={rowCount}
|
||||
filterMode="server"
|
||||
columnFilters={columnFilters}
|
||||
|
|
|
|||
|
|
@ -85,13 +85,19 @@ describe("row action cells", () => {
|
|||
expect(deps.onKeyHashClick).toHaveBeenCalledWith("sk-hash-9");
|
||||
});
|
||||
|
||||
it("reports the session id from the session cell", async () => {
|
||||
it("reports the clicked row from the session cell, so two rows sharing a session id stay distinguishable", async () => {
|
||||
const user = userEvent.setup();
|
||||
const deps = { onKeyHashClick: vi.fn(), onSessionClick: vi.fn() };
|
||||
renderRows([logEntry({ request_id: "req-sess", session_id: "sess-42" })], deps);
|
||||
renderRows(
|
||||
[
|
||||
logEntry({ request_id: "req-key-a", session_id: "sess-42", api_key: "key-a" }),
|
||||
logEntry({ request_id: "req-key-b", session_id: "sess-42", api_key: "key-b" }),
|
||||
],
|
||||
deps,
|
||||
);
|
||||
|
||||
await user.click(screen.getByText("sess-42"));
|
||||
expect(deps.onSessionClick).toHaveBeenCalledWith("sess-42");
|
||||
await user.click(screen.getAllByText("sess-42")[1]);
|
||||
expect(deps.onSessionClick).toHaveBeenCalledWith(expect.objectContaining({ request_id: "req-key-b" }));
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import { AgentBadge, AgentIcon, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } fr
|
|||
|
||||
export interface RequestLogsTableColumnsDeps {
|
||||
onKeyHashClick: (keyHash: string) => void;
|
||||
onSessionClick: (sessionId: string) => void;
|
||||
onSessionClick: (log: LogEntry) => void;
|
||||
}
|
||||
|
||||
const readMetaString = (metadata: Record<string, unknown> | undefined, key: string): string | undefined => {
|
||||
|
|
@ -61,7 +61,7 @@ export const getRequestLogsTableColumns = ({
|
|||
const isAgent = AGENT_CALL_TYPES.includes(log.call_type);
|
||||
const sessionLlmCount = log.session_llm_count ?? (isMcp || isAgent ? 0 : sessionCount);
|
||||
const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0);
|
||||
const sessionMcpCount = log.session_mcp_count ?? (isMcp ? sessionCount : 0);
|
||||
const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0);
|
||||
|
||||
if (isMcp) return <McpBadge />;
|
||||
if (isAgent && sessionCount <= 1) return <AgentBadge />;
|
||||
|
|
@ -113,7 +113,7 @@ export const getRequestLogsTableColumns = ({
|
|||
header: "Session ID",
|
||||
size: 120,
|
||||
enableSorting: false,
|
||||
cell: ({ row }) => <IdCell value={row.original.session_id} onClick={onSessionClick} />,
|
||||
cell: ({ row }) => <IdCell value={row.original.session_id} onClick={() => onSessionClick(row.original)} />,
|
||||
},
|
||||
{
|
||||
id: "request_id",
|
||||
|
|
|
|||
|
|
@ -47,6 +47,5 @@ export type LogEntry = {
|
|||
mcp_tool_call_count?: number;
|
||||
mcp_tool_call_spend?: number;
|
||||
session_llm_count?: number;
|
||||
session_mcp_count?: number;
|
||||
session_agent_count?: number;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -12,6 +12,9 @@ export const ERROR_CODE_OPTIONS: { label: string; value: string }[] = [
|
|||
{ label: "529 - Overloaded", value: "529" },
|
||||
];
|
||||
|
||||
/** Page sizes the logs tables offer; the first entry is the default. */
|
||||
export const LOGS_PAGE_SIZE_OPTIONS = [10, 25, 50, 100];
|
||||
|
||||
/** Call types that represent MCP tool invocations (shared across columns, index, drawer). */
|
||||
export const MCP_CALL_TYPES = ["call_mcp_tool", "list_mcp_tools"];
|
||||
|
||||
|
|
|
|||
|
|
@ -181,6 +181,7 @@ export function useLogFilterLogic({
|
|||
sort_by: sortBy,
|
||||
sort_order: sortOrder,
|
||||
exclude_internal_health_checks: excludeInternalHealthChecks,
|
||||
group_by_session: true,
|
||||
},
|
||||
});
|
||||
},
|
||||
|
|
|
|||
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -56755,6 +56755,8 @@ export interface operations {
|
|||
sort_order?: string | null;
|
||||
/** @description Exclude LiteLLM internal health check requests from results */
|
||||
exclude_internal_health_checks?: boolean;
|
||||
/** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */
|
||||
group_by_session?: boolean;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
|
|
@ -56867,6 +56869,8 @@ export interface operations {
|
|||
sort_order?: string | null;
|
||||
/** @description Exclude LiteLLM internal health check requests from results */
|
||||
exclude_internal_health_checks?: boolean;
|
||||
/** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */
|
||||
group_by_session?: boolean;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue