From 8c646a916c8860a04442645973974f71ee9c1c0b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:19:17 -0700 Subject: [PATCH 1/6] feat(cli): pre-fill the SSO verification code in the browser when the proxy allows it --- litellm/proxy/client/cli/commands/auth.py | 22 ++++-- .../proxy/client/cli/test_auth_commands.py | 70 +++++++++++++++++++ type-discipline-budget.json | 6 +- 3 files changed, 90 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 2fad9f933c1..a96d3fb9c85 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -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 diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 1d0a99b8e0a..821323e722c 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -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() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..2836b22c3d0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22364 }, "LIT002": { - "limit": 26777 + "limit": 26776 }, "LIT003": { "limit": 269 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16507 + "limit": 16506 }, "LIT011": { "limit": 5535 }, "LIT012": { - "limit": 4495 + "limit": 4492 } } From 8e3566d2f7564438c8c33ab66ce831aee416059e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:38:43 -0700 Subject: [PATCH 2/6] fix(guardrails): forward mode and streaming params to crowdstrike_aidr handler (#39317) * fix(guardrails): forward mode and streaming params to crowdstrike_aidr handler Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): drop stream_chunk_builder patch from crowdstrike cadence test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): type test params and cover unsupported crowdstrike mode rejection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crowdstrike_aidr/__init__.py | 13 +- .../crowdstrike_aidr/crowdstrike_aidr.py | 40 ++++- .../guardrail_hooks/crowdstrike_aidr.py | 13 +- .../guardrail_hooks/test_crowdstrike_aidr.py | 142 +++++++++++++++++- 4 files changed, 197 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py index 5e75b7d4d94..c88e6e97a96 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index c8284fac440..f7f500b1adc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -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: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index f47c38af3e3..6beca030a3a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py @@ -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]): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index ec7854b9a35..1b50ea53db2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -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 From 8cbaba88635871b44096488d6264490eac700b9f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:39:36 -0700 Subject: [PATCH 3/6] fix(mcp): gate the connect-time OBO pre-flight on the key's allowed servers (#39447) The pre-flight token exchange resolved its target from the requested path, so a key not entitled to that server still drove an outbound exchange at the IdP and populated the credential cache before the later access check denied it. Resolve the target through the same allowed-server set the rest of the route uses and skip the exchange when the requested server is not in it. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/server.py | 7 ++ .../mcp_server/test_mcp_server.py | 73 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 989b08b929a..af3ff6714af 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 82f74cda835..3f6d8f8837c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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.""" From c19d49d9192f26a7e15c9571d1c3a4e7ee45092a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:40:13 -0700 Subject: [PATCH 4/6] fix(responses): keep provider response headers in streaming logging callbacks (#38131) * fix(responses): keep provider response headers in streaming logging callbacks The responses streaming iterator captures the provider's HTTP response headers into its own _hidden_params, but never puts them on the completed response, and the model_validate(model_dump()) copy made for logging drops pydantic private attributes. Success callbacks and StandardLoggingPayload.hidden_params.additional_headers therefore saw an empty dict for streaming /v1/responses, so Azure's apim-request-id was unreadable from the callback payload. Restore the headers on the nested response of the logging copy, preferring any the provider transform already set (the fake_stream path) and falling back to the ones the iterator captured from the stream. Skipped when the copy fell back to the original event, so a serialization failure never leaves logging-only state on the caller's object. * fix: satisfy LIT002 mutable-collection gate in header restore --------- Co-authored-by: Yucheng Zhu --- litellm/responses/streaming_iterator.py | 40 +++++ .../responses/test_streaming_iterator.py | 161 +++++++++++++++++- 2 files changed, 200 insertions(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7871c85220c..f271655f5e3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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""" diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 9edcaaef034..c226c0b4d09 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -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 == {} From 44a6c659fb831bc3da929f020f2cd4f3e508e384 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 2 Sep 2026 17:45:58 -0700 Subject: [PATCH 5/6] fix(mcp): fence an outbound-token write against an overlapping invalidation (#35398) get_or_compute single-flights concurrent misses under a per-key lock, but invalidate() deletes outside it. A mint already awaiting the IdP when an invalidation ran wrote its result into the slot afterwards, so a bearer minted before the invalidation was served after it for its full TTL; the upstream-401 retry then re-presented the bearer the server had just rejected. The cache now keeps a per-key generation beside the single-flight lock. invalidate() bumps it and get_or_compute snapshots it inside the lock right before minting, skipping the cache write when it no longer matches. The token is still returned to the caller it was minted for, so only the caching is fenced. Generation and lock share one object that a running mint binds to a local, which pins the weak map's entry for exactly as long as a write can still land; a generation held weakly on its own would be collected across the await and the re-check would silently pass. --- .../outbound_credentials/token_endpoint.py | 67 ++++++++--- .../test_token_endpoint.py | 105 ++++++++++++++++++ 2 files changed, 158 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py index 84f714db449..6f1e5baa109 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py @@ -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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py index 5f277db2f72..db3a1a386a3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -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() From ff1f21aea92574b3df1b9a00b2cdd70fa5e17dda Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 17:54:17 -0700 Subject: [PATCH 6/6] fix(ui): paginate request logs by session groups server-side (#39257) * fix(ui): paginate request logs by session groups server-side The logs table server-paginated raw spend logs and then collapsed multi-call sessions client-side, so a page could render 3 rows while the footer claimed 25 and sessions straddled pages. Adds an opt-in group_by_session param to /spend/logs/ui that pages and counts one representative row per session (DISTINCT ON, newest non-MCP call), keeps the bounded count contract, enriches whole-session llm/agent composition counts, and deletes the client-side collapse pipeline. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve * feat(ui): add a 10 rows-per-page option and default request logs to it Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve * fix(ui): key session aggregates per api key in the logs enrichment Grouped pagination splits a reused session id into one row per api key, but the enrichment still aggregated by session_id alone, so both rows showed combined spend and counts. The aggregate query now groups by (session_id, api_key), the count folds into it (the separate group_by query is deleted), and each row reads its own key's totals. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve * fix(ui): treat an empty api_key as a real session group value The spend-log schema defaults api_key to an empty string; truthiness guards in the enrichment treated it as missing, so keyless multi-call sessions lost their count and spend. Only None means missing now. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve --------- Co-authored-by: Claude Fable 5 --- .../spend_management_endpoints.py | 118 ++++++++------ tests/e2e/ui/helpers/traffic.ts | 3 + .../e2e/ui/tests/logs/logsPagination.spec.ts | 150 ++++++++++++++++++ .../test_spend_management_endpoints.py | 145 +++++++++++++++-- .../test_spend_query_optimization.py | 106 +++++++++++++ .../src/components/networking.tsx | 1 + .../view_logs/RequestLogsPanel.test.tsx | 45 +++--- .../components/view_logs/RequestLogsPanel.tsx | 63 +------- .../components/view_logs/RequestLogsTable.tsx | 4 +- .../RequestLogsTableColumns.test.tsx | 14 +- .../view_logs/RequestLogsTableColumns.tsx | 6 +- .../src/components/view_logs/columns.tsx | 1 - .../src/components/view_logs/constants.ts | 3 + .../components/view_logs/log_filter_logic.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 15 files changed, 512 insertions(+), 152 deletions(-) create mode 100644 tests/e2e/ui/tests/logs/logsPagination.spec.ts diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c48750cea72..bb7dfafb297 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -55,6 +55,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), @@ -144,21 +148,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): @@ -242,18 +241,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}) @@ -2290,6 +2277,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. @@ -2644,12 +2635,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 """ @@ -2660,21 +2655,36 @@ 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, cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_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 {_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 {_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) @@ -4075,11 +4085,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. @@ -4101,7 +4112,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( { @@ -4110,15 +4120,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 @@ -4130,38 +4133,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( @@ -4174,14 +4185,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: diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index 25eb671fd0e..7f8417cdffb 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -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); diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts new file mode 100644 index 00000000000..416e1c3171f --- /dev/null +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -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 { + await navigateToPage(page, Page.Logs); + await dismissFeedbackPopup(page); + await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 }); +} + +async function openFilterDrawer(page: PlaywrightPage): Promise { + 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 { + 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 { + 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(); + }); +}); diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index a0dcbf802ef..30b086bab61 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3959,18 +3959,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, } ] ) @@ -3995,6 +3995,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 @@ -4003,13 +4005,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(): @@ -4033,14 +4148,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, @@ -4089,13 +4203,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, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index ef68d9ce178..27e633099f0 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -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 diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index db7baf5e171..63232afc46b 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2042,6 +2042,7 @@ interface UiSpendLogsParams { min_spend?: number; max_spend?: number; exclude_internal_health_checks?: boolean; + group_by_session?: boolean; } interface UiSpendLogsCallOptions { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index 22e3f635b50..c368a43dd1e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -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("closing the drawer removes ?log_id= from the URL and closes the drawer", async () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 52ea78abf5e..2ebabf64ab0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -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" }; interface RequestLogsPanelProps { @@ -35,12 +35,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({ pageIndex: 0, pageSize: PAGE_SIZE }); const [sorting, setSorting] = useState(DEFAULT_LOGS_SORTING); @@ -157,49 +151,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const isDrawerOpen = displayLog !== null || displaySessionId !== null; - const rows = useMemo(() => { - const searchedLogs = filteredLogs.data; - - const sessionCompositionById = searchedLogs.reduce>((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(); - 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); @@ -258,13 +210,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( diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx index 7146cf33847..4159b3b699b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx @@ -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} diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 3d0ea03c5d5..ce59c62f1c4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -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" })); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index cf776515bd4..b9058d02a6b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -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 | 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 ; if (isAgent && sessionCount <= 1) return ; @@ -113,7 +113,7 @@ export const getRequestLogsTableColumns = ({ header: "Session ID", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => onSessionClick(row.original)} />, }, { id: "request_id", diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index eef957922d7..2f3a3681352 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -46,6 +46,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; }; diff --git a/ui/litellm-dashboard/src/components/view_logs/constants.ts b/ui/litellm-dashboard/src/components/view_logs/constants.ts index 5b0b1d0fee3..1c17f398e35 100644 --- a/ui/litellm-dashboard/src/components/view_logs/constants.ts +++ b/ui/litellm-dashboard/src/components/view_logs/constants.ts @@ -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"]; diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index 3b8d96596de..acd5be06593 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -181,6 +181,7 @@ export function useLogFilterLogic({ sort_by: sortBy, sort_order: sortOrder, exclude_internal_health_checks: excludeInternalHealthChecks, + group_by_session: true, }, }); }, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index be81450b941..7583b6f6df1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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;