diff --git a/deploy/lens/README.md b/deploy/lens/README.md index 4971bfae671..6f8c7d3940e 100644 --- a/deploy/lens/README.md +++ b/deploy/lens/README.md @@ -2,62 +2,74 @@ Lens reviews recorded activity and saves evidence-linked findings in the LiteLLM dashboard under Observability, Lens (`/ui/lens/`) -## Install the release stack +## Install -Each stable, RC, and dev release containing Lens publishes the worker at the same version on GHCR and Docker Hub. Use the [LiteLLM releases page](https://github.com/BerriAI/litellm/releases) to select a version that includes the coordinated worker release +Build LiteLLM and its worker from the same source commit with the same release identity. The worker runs separately and connects to your gateway using a limited worker token -For a new local installation, install Docker with Compose, download the two release files, and create a private environment file. Replace `X.Y.Z` with the release version, without `v` (RCs use `X.Y.Z-rc.N`) +### New local installation + +Install Docker with Compose and Git. This builds LiteLLM and its worker from the same checkout and starts the existing local tracing stack: ```bash -mkdir litellm-lens -cd litellm-lens -LENS_RELEASE=X.Y.Z -curl -fSLo compose.yaml "https://raw.githubusercontent.com/BerriAI/litellm/v${LENS_RELEASE}/deploy/lens/stack.yaml" -curl -fSLo config.yaml "https://raw.githubusercontent.com/BerriAI/litellm/v${LENS_RELEASE}/deploy/lens/config.yaml" -umask 077 -printf 'LITELLM_VERSION=%s\nLITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' \ - "$LENS_RELEASE" "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env -printf 'POSTGRES_PASSWORD=%s\nCLICKHOUSE_PASSWORD=%s\n' \ - "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" >> .env -docker compose up -d +git clone https://github.com/BerriAI/litellm.git +cd litellm +export LITELLM_RELEASE_TAG="sha-$(git rev-parse HEAD)" +export LENS_WORKER_IMAGE="litellm-lens-worker:${LITELLM_RELEASE_TAG}" +export OPENAI_API_KEY='sk-...' +docker build --build-arg LITELLM_RELEASE_TAG="$LITELLM_RELEASE_TAG" \ + -f deploy/lens/Dockerfile -t "$LENS_WORKER_IMAGE" . +docker compose -f docker/docker-compose.tracing.yml up -d --build ``` -Open `http://localhost:4000/ui/`, log in as `admin` with `LITELLM_MASTER_KEY` from `.env`, and add a model in the dashboard. In Lens, select **Connect worker**, choose that model and a monthly budget, then **Get install command**. Expand **Using Docker Compose or Helm?**, copy the worker token, and add `LENS_WORKER_TOKEN=` to `.env` +Open `http://localhost:4002/ui/` and sign in as `admin` with password `sk-1234`. Go to **Lens > Investigations > Connect worker**, choose a model and monthly budget, then **Get install command**. Expand **Using Docker Compose or Helm?** and copy the worker token. In the same terminal, run: ```bash -docker compose --profile lens up -d +export LITELLM_URL=http://litellm:4000 +export LENS_WORKER_TOKEN='' +docker compose -f docker/docker-compose.tracing.yml -f deploy/lens/compose.yaml up -d ``` -The stack starts LiteLLM, PostgreSQL, ClickHouse, and the worker from published images. The dashboard shows **Worker connected**. The worker has a limited token, no database credentials, and no provider keys. The stack exposes only the dashboard on localhost; use your normal ingress and managed databases for a public production deployment +The worker joins the gateway's Docker network, and the dashboard shows **Worker connected**. Save the token privately for restarts and upgrades -Keep `.env` private and preserve its salt key. Keep both named database volumes. To upgrade, wait for active investigations to finish, stop the worker, change only `LITELLM_VERSION`, then pull and recreate the stack: +This stack is for local evaluation: it binds to localhost and uses development database credentials. For a hosted deployment, keep your normal database, keys, networking, and deployment process. Build both images from one source revision with the same `LITELLM_RELEASE_TAG`, publish the worker to your registry, and set `LENS_WORKER_IMAGE` on LiteLLM to that image + +### Existing LiteLLM installation + +Keep your deployment and PostgreSQL database. A working gateway/worker pair can stay as it is until you upgrade both. For a gateway built from source, use its exact commit and `LITELLM_RELEASE_TAG`; a release version or the latest commit on `main` is not a substitute for that source identity + +The public development package is `ghcr.io/berriai/litellm-lens-worker-dev:sha-`. It publishes amd64 images on Lens-related changes, so an arbitrary source commit may have no image. Check the exact image exists before using it. If it is unavailable, your gateway uses a different release identity, or you need native arm64, build the worker from the gateway's checkout: ```bash -docker compose --profile lens stop lens-worker -# Update LITELLM_VERSION in .env to the new release -docker compose --profile lens pull -docker compose --profile lens up -d +export LITELLM_RELEASE_TAG='' +export LENS_WORKER_IMAGE='/litellm-lens-worker:' +docker build --build-arg LITELLM_RELEASE_TAG="$LITELLM_RELEASE_TAG" \ + -f deploy/lens/Dockerfile -t "$LENS_WORKER_IMAGE" . ``` -This preserves your investigations, findings, model credentials, and worker token. Never use `down -v` during an upgrade. If moving from an existing installation, keep its databases and add the standalone worker instead of creating an empty replacement stack +For a remote worker host, publish that image to a registry the host can pull from. Set the gateway's `LENS_WORKER_IMAGE` to the resulting image reference, restart the gateway using its normal deployment process, then copy its install command. Prefer the published image digest for hosted installations. Do not change the gateway's release identity just to accept another worker + +For Kubernetes or Render, run the standalone worker using `LITELLM_URL` and `LENS_WORKER_TOKEN` from setup. Keep existing databases and secrets. The worker needs no inbound port. ## Helm -The componentized `helm/litellm` chart includes an optional Lens worker. Configure PostgreSQL and ClickHouse as usual, install the chart, then obtain a limited worker token from Lens setup. Store it in a Kubernetes Secret and enable the worker in your values: +The componentized source chart at `helm/litellm` includes an optional Lens worker. Use the chart from the same checkout as your gateway and keep your component image overrides in your values. Configure PostgreSQL and ClickHouse as usual, install the chart, then obtain a limited worker token from Lens setup. Store it in a Kubernetes Secret and enable the worker in your values: ```yaml lensWorker: enabled: true + image: + repository: + digest: sha256: tokenSecret: name: litellm-lens-worker key: token ``` -Published release charts pin the worker's approved image digest. Source charts without a digest default to the chart's application version. The chart connects the worker to the backend service. Keep these values and the Secret when upgrading the chart so the gateway and worker upgrade together. `lensWorker.replicaCount` controls simultaneous investigations. To use a private registry or external proxy, set `lensWorker.image.repository`, `lensWorker.image.digest` (or `tag` for a source build), and `lensWorker.url`. A digest takes precedence over the tag. The dashboard uses the chart's worker image for standalone install commands too +Set the worker repository and digest explicitly to an image built from the gateway's source commit and release identity. The chart connects the worker to the backend service. Keep these values and the Secret when upgrading the chart and update the gateway and worker image overrides together. `lensWorker.replicaCount` controls simultaneous investigations. To use a private registry or external proxy, set `lensWorker.image.repository`, `lensWorker.image.digest` (or `tag` for a source build), and `lensWorker.url`. A digest takes precedence over the tag. The dashboard uses the chart's worker image for standalone install commands too ## Standalone worker -Upgrade your existing LiteLLM proxy to a release that includes Lens with PostgreSQL and agent tracing. Configure one ClickHouse URL for trace writes, bounded reads, and Lens queries: +Start with a source deployment that includes Lens, PostgreSQL, and agent tracing, and prepare its matching worker as described above. Configure one ClickHouse URL for trace writes, bounded reads, and Lens queries: ```yaml general_settings: @@ -74,13 +86,13 @@ Retention changes require a proxy restart. ClickHouse removes expired rows durin In **Lens > Investigations**, click **Connect worker**, choose an analysis model and monthly limit, then **Get install command**. Use **Advanced options** to select an existing virtual key or change the proxy URL if the server running Docker needs a different network address. Copy the command and run it on your server. The dashboard shows **Worker connected** when the container checks in -The command already contains the compatible worker image and one worker token. The selected virtual key stays on the proxy; its secret is never sent to the worker. No source checkout, environment file, or second LiteLLM deployment is needed. Keep the command private because it includes the token. The LiteLLM release provides the dashboard and APIs; the container only runs background analysis +The command already contains the compatible worker image and one worker token. The selected virtual key stays on the proxy; its secret is never sent to the worker. Once the matching image is available on the worker host, no second LiteLLM deployment is needed. Keep the command private because it includes the token. The LiteLLM release provides the dashboard and APIs; the container only runs background analysis -The dashboard selects the worker image matching the running gateway release. Release images support Linux amd64 and arm64. CI also publishes `:sha-` development images; use those only with a gateway built from the same commit and release tag +The dashboard uses the gateway's `LENS_WORKER_IMAGE` override when set. Public `:sha-` development images must match both the gateway commit and release identity. Build from source for the worker host's native architecture After upgrading the gateway, update the worker image and redeploy it while keeping its proxy URL and token. Existing containers do not update automatically. If an investigation reports a worker compatibility error, update the image before retrying -For deployments managed with Compose, download `compose.yaml` and provide `LITELLM_URL`, `LENS_WORKER_TOKEN`, and `LITELLM_VERSION` (without `v`) in a private environment file. To use another registry, set `LENS_WORKER_IMAGE` to the compatible image instead of setting a version: +For deployments managed with Compose, download `compose.yaml` and provide `LITELLM_URL`, `LENS_WORKER_TOKEN`, and an explicit `LENS_WORKER_IMAGE` in a private environment file: ```bash docker compose --env-file /path/to/lens.env -f compose.yaml up -d @@ -219,9 +231,9 @@ Upgrades using `--use_prisma_db_push` stop before schema changes if any legacy L ## Release compatibility -Released gateway and worker images carry `LITELLM_RELEASE_TAG`. A worker announces its release and protocol before claiming an investigation. A mismatch returns HTTP 409 with the required image, leaving queued investigations untouched. During a rolling upgrade, workers wait for a gateway from their release +Gateway and worker builds carry the same `LITELLM_RELEASE_TAG`. A worker announces its release and protocol before claiming an investigation. A mismatch returns HTTP 409 with the required image, leaving queued investigations untouched. During a rolling upgrade, workers wait for a gateway from their release -The dashboard reads its image from the running gateway. `LENS_WORKER_IMAGE` overrides the registry/image for private deployments. Worker-only Compose accepts `LITELLM_VERSION` (without `v`) or an explicit `LENS_WORKER_IMAGE`. Release workers are available as `ghcr.io/berriai/litellm-lens-worker:vX.Y.Z` and `docker.io/litellm/litellm-lens-worker:vX.Y.Z`, including matching RC/dev suffixes, on amd64 and arm64 +The dashboard reads its image from the running gateway. `LENS_WORKER_IMAGE` overrides the registry/image for private deployments. Set an explicit `LENS_WORKER_IMAGE` for worker-only Compose. Verify that the image exists and matches the gateway before deploying it For source development, use `make lens-dev`, which gives the proxy and source worker the same commit identity. For custom containers, build both from the same checkout with `--build-arg LITELLM_RELEASE_TAG=sha-$(git rev-parse HEAD)` and set the proxy's `LENS_WORKER_IMAGE` to the worker image you built. An unlabelled custom build refuses worker setup and claims instead of guessing from the Python package version. Normal package-index installations use their installed release version diff --git a/deploy/lens/compose.yaml b/deploy/lens/compose.yaml index 799f0a4fb1e..aa915fef663 100644 --- a/deploy/lens/compose.yaml +++ b/deploy/lens/compose.yaml @@ -1,6 +1,6 @@ services: lens-worker: - image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker:v${LITELLM_VERSION:?Set LITELLM_VERSION to the gateway release, without the v prefix}} + image: ${LENS_WORKER_IMAGE:-${LITELLM_VERSION:+ghcr.io/berriai/litellm-lens-worker:v}${LITELLM_VERSION:-}} environment: LITELLM_URL: ${LITELLM_URL:?Set the URL reachable from this container} LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:?Create a worker credential in the Lens UI} diff --git a/docker/docker-compose.tracing.yml b/docker/docker-compose.tracing.yml index 8f960d50872..c8d90fbc0ae 100644 --- a/docker/docker-compose.tracing.yml +++ b/docker/docker-compose.tracing.yml @@ -5,6 +5,8 @@ services: build: context: .. target: runtime + args: + LITELLM_RELEASE_TAG: ${LITELLM_RELEASE_TAG:-} command: ["--config", "/app/tracing-config.yaml", "--port", "4000"] environment: LITELLM_MASTER_KEY: sk-1234 @@ -15,6 +17,7 @@ services: CLICKHOUSE_URL: http://default:local-tracing@clickhouse:8123 CLICKHOUSE_DATABASE: litellm OPENAI_API_KEY: ${OPENAI_API_KEY:-} + LENS_WORKER_IMAGE: ${LENS_WORKER_IMAGE:-} volumes: - ./tracing-config.yaml:/app/tracing-config.yaml:ro ports: diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 245244250ee..1fd292b8137 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -1,3 +1,4 @@ +import functools import glob import os import random @@ -10,7 +11,8 @@ import time from collections.abc import Callable from dataclasses import dataclass, replace from pathlib import Path -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Final, Optional, Union +from urllib.parse import unquote, urlsplit from litellm_proxy_extras import prisma_toolchain from litellm_proxy_extras._logging import logger @@ -202,6 +204,66 @@ def _max_migration_timestamp(names) -> int: return max(_migration_timestamp(n) for n in names) +_REDACTED: Final = "REDACTED" +_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword")) + + +@functools.cache +def _secret_shape_redactor() -> Callable[[str], str]: + try: + from litellm._logging import redact_secrets + except ImportError: + return lambda text: text + return redact_secrets + + +def _url_passwords(url: str) -> frozenset[str]: + try: + parts: Final = urlsplit(url) + except ValueError: + return frozenset() + query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&")) + raw_query_passwords: Final = tuple( + value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS + ) + raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords + return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password) + + +def _configured_database_passwords() -> frozenset[str]: + database_url: Final = os.getenv("DATABASE_URL") + direct_url: Final = os.getenv("DIRECT_URL") + database_passwords: Final = _url_passwords(database_url) if database_url else frozenset() + direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset() + return database_passwords | direct_passwords + + +def _redact_credentials(text: str) -> str: + """Mask configured database passwords before passing the text to LiteLLM redaction.""" + passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True) + alternation: Final = "|".join(re.escape(password) for password in passwords) + password_pattern: Final = ( + re.compile(rf"(?P:|password=)(?:{alternation})(?=@|&|$|[\s'\"\]),])", re.IGNORECASE) + if passwords + else None + ) + result: Final = password_pattern.sub(rf"\g{_REDACTED}", text) if password_pattern is not None else text + return _secret_shape_redactor()(result) + + +def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]: + if isinstance(command, tuple): + return tuple(_redact_credentials(str(argument)) for argument in command) + if isinstance(command, list): + return [_redact_credentials(str(argument)) for argument in command] + return _redact_credentials(str(command)) + + +def _redact_command_error(error: subprocess.CalledProcessError) -> str: + redacted_command: Final = _redacted_command(error.cmd) + return str(subprocess.CalledProcessError(error.returncode, redacted_command)) + + def _get_prisma_command() -> str: """Get the Prisma command to use, bypassing Python wrapper in offline mode.""" if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): @@ -315,7 +377,8 @@ class ProxyExtrasDBManager: return False except subprocess.CalledProcessError as e: logger.warning( - f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}" + f"Error creating baseline migration: {_redact_command_error(e)}, " + f"{_redact_credentials(str(e.stderr))}, {_redact_credentials(str(e.stdout))}" ) raise e @@ -1572,6 +1635,11 @@ class ProxyExtrasDBManager: f"Error: {stderr}" ) raise + else: + logger.error( + "prisma migrate deploy failed with an error the resolver does not handle: " + f"{_redact_credentials(stderr)}" + ) else: if ProxyExtrasDBManager.spend_logs_is_partitioned(): raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR) @@ -1586,7 +1654,7 @@ class ProxyExtrasDBManager: ) return True except subprocess.TimeoutExpired: - logger.warning( + logger.error( "Attempt %s timed out. Raise %s if this database needs longer to apply its schema.", attempt + 1, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR, @@ -1599,7 +1667,12 @@ class ProxyExtrasDBManager: if attempts_left > 0 else "" ) - logger.info(f"The process failed to execute. Details: {e}.{retry_msg}") + stderr_detail: Final = ( + f" stderr: {_redact_credentials(str(e.stderr))}" if e.stderr else "" + ) + logger.error( + f"The process failed to execute. Details: {_redact_command_error(e)}.{stderr_detail}{retry_msg}" + ) time.sleep(random.randrange(5, 15)) finally: os.chdir(original_dir) diff --git a/litellm/__init__.py b/litellm/__init__.py index b6428b51bfb..fea7a27a5fb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1474,6 +1474,7 @@ from .rust_bridge import rust from .rag.main import * from .sandbox.main import * from .decisions.main import * +from .tool_loop import ToolLoopMaxRoundsExceeded, arun_tool_loop, run_tool_loop from .search.main import * from .realtime_api.main import ( _arealtime, diff --git a/litellm/constants.py b/litellm/constants.py index 80e4c9091ba..f4b9219b910 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2235,3 +2235,5 @@ HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset( ".ruff_cache", } ) + +DEFAULT_TOOL_LOOP_MAX_ROUNDS: Final = 20 diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py new file mode 100644 index 00000000000..64d6de51818 --- /dev/null +++ b/litellm/tool_loop.py @@ -0,0 +1,172 @@ +"""Client-side tool-calling loop helpers for litellm.completion.""" + +from collections.abc import Awaitable, Callable, Sequence +from typing import TypeAlias, cast + +from typing_extensions import TypedDict, Unpack + +from litellm.constants import DEFAULT_TOOL_LOOP_MAX_ROUNDS +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, + ChatCompletionToolCallFunctionChunk, + ChatCompletionToolMessage, + ChatCompletionToolParam, +) +from litellm.types.utils import ChatCompletionMessageToolCall, Message, ModelResponse + +ToolExecutor: TypeAlias = Callable[[ChatCompletionMessageToolCall], ChatCompletionToolMessage] +AsyncToolExecutor: TypeAlias = Callable[[ChatCompletionMessageToolCall], Awaitable[ChatCompletionToolMessage]] + + +class _ToolLoopCompletionKwargs(TypedDict, total=False, extra_items=object): + """Extra keywords forwarded verbatim to ``litellm.completion``, which owns their contract.""" + + +class ToolLoopMaxRoundsExceeded(RuntimeError): + max_rounds: int + + def __init__(self, max_rounds: int) -> None: + self.max_rounds = max_rounds + super().__init__(f"model still requested tool calls on round {max_rounds} of {max_rounds}") + + +def _validate_tool_loop_args(max_rounds: int, completion_kwargs: _ToolLoopCompletionKwargs) -> None: + if max_rounds < 1: + raise ValueError(f"max_rounds must be >= 1, got {max_rounds}") + if completion_kwargs.get("stream"): + raise ValueError("run_tool_loop requires whole responses; stream=True is not supported") + + +def _expect_model_response(response: object) -> ModelResponse: + if not isinstance(response, ModelResponse): + raise TypeError(f"run_tool_loop requires completion to return a ModelResponse, got {type(response).__name__}") + return response + + +def _assistant_tool_call(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionAssistantToolCall: + return ChatCompletionAssistantToolCall( + id=tool_call.id, + type="function", + function=ChatCompletionToolCallFunctionChunk( + name=tool_call.function.name, arguments=tool_call.function.arguments + ), + ) + + +def _assistant_message( + message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...] +) -> ChatCompletionAssistantMessage: + return cast( # cast-ok: dict literal with thinking_blocks and reasoning_items spread in only when set + "ChatCompletionAssistantMessage", + { + "role": "assistant", + "content": message.content, + "tool_calls": [_assistant_tool_call(tc) for tc in tool_calls], + **{ + key: value + for key, value in ( + ("thinking_blocks", getattr(message, "thinking_blocks", None)), + ("reasoning_items", getattr(message, "reasoning_items", None)), + ) + if value is not None + }, + }, + ) + + +def _function_tool_call(tool_call: object) -> ChatCompletionMessageToolCall: + if not isinstance(tool_call, ChatCompletionMessageToolCall): + raise TypeError( + f"run_tool_loop only executes function tool calls, got custom tool call {getattr(tool_call, 'id', None)}" + ) + return tool_call + + +def _function_tool_calls(message: Message) -> tuple[ChatCompletionMessageToolCall, ...]: + return tuple(_function_tool_call(tool_call) for tool_call in message.tool_calls or ()) + + +def run_tool_loop( + *, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam], + execute_tool: ToolExecutor, + max_rounds: int = DEFAULT_TOOL_LOOP_MAX_ROUNDS, + **completion_kwargs: Unpack[_ToolLoopCompletionKwargs], # kwargs-ok: forwarded verbatim to litellm.completion +) -> str | None: + """Call completion, execute each requested tool, and repeat until the model answers. + + Returns the final assistant message content. Raises ToolLoopMaxRoundsExceeded when the + model is still requesting tools after max_rounds completions. + + Example: + execute = functools.partial(run_repo_tool, repository="litellm", revision="main") + answer = litellm.run_tool_loop( + model="anthropic/claude-sonnet-5-5", messages=messages, tools=tools, execute_tool=execute + ) + """ + import litellm + + _validate_tool_loop_args(max_rounds, completion_kwargs) + history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns + for round_number in range(1, max_rounds + 1): + message = ( + _expect_model_response( + litellm.completion(model=model, messages=list(history), tools=list(tools), **completion_kwargs) + ) + .choices[0] + .message + ) + if not message.tool_calls: + return message.content + tool_calls = _function_tool_calls(message) + if round_number == max_rounds: + break + tool_results = tuple(execute_tool(tool_call) for tool_call in tool_calls) + history = (*history, _assistant_message(message, tool_calls), *tool_results) + raise ToolLoopMaxRoundsExceeded(max_rounds) + + +async def arun_tool_loop( + *, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam], + execute_tool: AsyncToolExecutor, + max_rounds: int = DEFAULT_TOOL_LOOP_MAX_ROUNDS, + **completion_kwargs: Unpack[_ToolLoopCompletionKwargs], # kwargs-ok: forwarded verbatim to litellm.acompletion +) -> str | None: + """Async version of run_tool_loop, awaiting each execute_tool call in order. + + Returns the final assistant message content. Raises ToolLoopMaxRoundsExceeded when the + model is still requesting tools after max_rounds completions. + + Example: + execute = functools.partial(arun_repo_tool, repository="litellm", revision="main") + answer = await litellm.arun_tool_loop( + model="anthropic/claude-sonnet-5-5", messages=messages, tools=tools, execute_tool=execute + ) + """ + import litellm + + _validate_tool_loop_args(max_rounds, completion_kwargs) + history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns + for round_number in range(1, max_rounds + 1): + message = ( + _expect_model_response( + await litellm.acompletion(model=model, messages=list(history), tools=list(tools), **completion_kwargs) + ) + .choices[0] + .message + ) + if not message.tool_calls: + return message.content + tool_calls = _function_tool_calls(message) + if round_number == max_rounds: + break + tool_results = tuple([await execute_tool(tool_call) for tool_call in tool_calls]) + history = (*history, _assistant_message(message, tool_calls), *tool_results) + raise ToolLoopMaxRoundsExceeded(max_rounds) diff --git a/pyproject.toml b/pyproject.toml index d765bccb552..f04e66a04eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "pydantic-settings>=2.14.1,<3.0", "jsonschema>=4.0.0,<5.0", "boto3>=1.43.1,<2.0", + "typing-extensions>=4.13.0,<5.0", ] [project.urls] diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index 70c49c5c256..a62af1b3725 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -177,3 +177,4 @@ hypothesis: >=6.165.10 # MPL 2.0 license pytest-rerunfailures: >=15.1 # MPL 2.0 license pytest-recording: >=0.13.4 # MIT license expression: >=5.6.0 # MIT License - https://github.com/cognitedata/Expression/blob/main/LICENSE +typing-extensions: >=4.13.0 # PSF-2.0 license - https://github.com/python/typing_extensions/blob/main/LICENSE diff --git a/tests/e2e/migrations/lens_compose_smoke.sh b/tests/e2e/migrations/lens_compose_smoke.sh index 69a13b58d88..6a57926895b 100644 --- a/tests/e2e/migrations/lens_compose_smoke.sh +++ b/tests/e2e/migrations/lens_compose_smoke.sh @@ -1,6 +1,19 @@ #!/usr/bin/env bash set -euo pipefail +worker_image() { + env -u LENS_WORKER_IMAGE -u LITELLM_VERSION \ + LITELLM_URL=http://litellm:4000 LENS_WORKER_TOKEN=config-test "$@" \ + docker compose --env-file /dev/null -f deploy/lens/compose.yaml config --images +} +[[ "$(worker_image LENS_WORKER_IMAGE=registry.example/lens:source)" == registry.example/lens:source ]] +[[ "$(worker_image LITELLM_VERSION=1.2.3)" == ghcr.io/berriai/litellm-lens-worker:v1.2.3 ]] +[[ "$(worker_image LENS_WORKER_IMAGE=registry.example/lens:source LITELLM_VERSION=1.2.3)" == registry.example/lens:source ]] +if worker_image > /dev/null 2>&1; then + printf 'Worker Compose accepted neither an image nor a release version\n' >&2 + exit 1 +fi + qa_dir=$(mktemp -d) master_key="sk-$(openssl rand -hex 32)" compose=(docker compose -p lens-compose-ci --env-file "$qa_dir/env" -f deploy/lens/stack.yaml) diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index bce09b49e10..d44dfc625aa 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -99,7 +99,7 @@ export const MIGRATED_E2E_PAGES: Readonly> = { group: "Settings", content: { role: "heading", name: "UI Theme Customization" }, }, - logs: { segment: "logs", linkName: "Logs", content: { role: "heading", name: "Request Logs" } }, + logs: { segment: "logs", linkName: "Logs", content: { role: "tab", name: "Request Logs" } }, "admin-panel": { segment: "admin-panel", linkName: "Admin Settings", diff --git a/tests/integration/_support/bedrock_runtime_peer.py b/tests/integration/_support/bedrock_runtime_peer.py index 3a547260590..64115d9eaa0 100644 --- a/tests/integration/_support/bedrock_runtime_peer.py +++ b/tests/integration/_support/bedrock_runtime_peer.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import json import re import threading diff --git a/tests/integration/database/test_v1_migration_error_logging.py b/tests/integration/database/test_v1_migration_error_logging.py new file mode 100644 index 00000000000..c3a3d8c43d6 --- /dev/null +++ b/tests/integration/database/test_v1_migration_error_logging.py @@ -0,0 +1,558 @@ +from __future__ import annotations + +import json +import os +import re +import signal +import socket +import socketserver +import subprocess +import sys +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager, suppress +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal, cast +from urllib.parse import quote, urlsplit, urlunsplit + +import psycopg +import pytest +from psycopg import sql +from psycopg.types.json import Jsonb + +from tests.integration._support.client import eventually +from tests.integration._support.process import _free_port as free_port + +REPO_ROOT: Final = Path(__file__).resolve().parents[3] +MIGRATIONS_DIR: Final = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" +MIGRATION_NAME: Final = "20260921190000_agent_identity" +pytestmark: Final = pytest.mark.timeout(300) +PASSWORD: Final = "wr ong'pw9" +FRAGMENTS: Final = ("wr ong", "wr+ong", "ong'pw9", "wr%20ong", "ong%27pw9", "pw9") +SHIPPED_MIGRATIONS: Final = tuple( + sorted(path.name for path in MIGRATIONS_DIR.iterdir() if path.is_dir() and path.name != "0_init") +) +LOG_PREFIX: Final = r"^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3} - [^\n]+ - " +LOG_RECORD_START: Final = rf"{LOG_PREFIX}(?:DEBUG|INFO|WARNING|ERROR|CRITICAL) - " +ERROR_RECORD: Final = re.compile(rf"(?ms)^({LOG_PREFIX}ERROR - .*?)(?={LOG_RECORD_START}|\Z)") +RETRY_COUNT: Final = re.compile(r"Retrying\.\.\. \((\d+) attempts left\)") +FRAGMENT_PATTERN: Final = re.compile( + "|".join(re.escape(fragment) for fragment in sorted(FRAGMENTS, key=len, reverse=True)) +) + + +@dataclass(frozen=True, slots=True) +class MigrationResult: + returncode: int + output: str + + +def _migration_log_path(test_name: str, tmp_path: Path) -> Path: + log_directory: Final = ( + Path(os.environ["INTEGRATION_RESULTS_DIR"]) if "INTEGRATION_RESULTS_DIR" in os.environ else tmp_path + ) + log_path: Final = log_directory / f"v1-migration-{test_name}-{uuid.uuid4().hex[:8]}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + return log_path + + +def _database_url( + admin_url: str, + role: str, + password: str, + database: str, + *, + encode_password: bool = True, +) -> str: + parsed: Final = urlsplit(admin_url) + authority: Final = parsed.netloc.rsplit("@", 1)[-1] + encoded_password: Final = quote(password, safe="") if encode_password else password + netloc: Final = f"{quote(role, safe='')}:{encoded_password}@{authority}" + return urlunsplit(parsed._replace(netloc=netloc, path=f"/{database}")) + + +def _replace_port(database_url: str, port: int, hostname: str | None = None) -> str: + parsed: Final = urlsplit(database_url) + target_host: Final = hostname or parsed.hostname + assert target_host is not None + host: Final = f"[{target_host}]" if ":" in target_host else target_host + userinfo: Final = parsed.netloc.rsplit("@", 1)[0] + return urlunsplit(parsed._replace(netloc=f"{userinfo}@{host}:{port}")) + + +def _unreachable_database_url(database_url: str) -> str: + parsed: Final = urlsplit(database_url) + url: Final = _database_url( + database_url, + parsed.username or "", + PASSWORD, + parsed.path.lstrip("/"), + encode_password=False, + ) + return _replace_port(url, free_port()) + + +@contextmanager +def owned_database(password: str) -> Iterator[str]: + admin_url: Final = os.environ["DATABASE_URL"] + role: Final = f"v1_migration_{uuid.uuid4().hex}" + database: Final = f"v1_migration_{uuid.uuid4().hex}" + database_url: Final = _database_url(admin_url, role, password, database) + try: + with psycopg.connect(admin_url, autocommit=True) as admin: + admin.execute( + sql.SQL("CREATE ROLE {} WITH LOGIN PASSWORD {}").format(sql.Identifier(role), sql.Literal(password)) + ) + admin.execute(sql.SQL("CREATE DATABASE {} OWNER {}").format(sql.Identifier(database), sql.Identifier(role))) + yield database_url + finally: + with psycopg.connect(admin_url, autocommit=True) as admin: + admin.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(sql.Identifier(database))) + admin.execute(sql.SQL("DROP ROLE IF EXISTS {}").format(sql.Identifier(role))) + + +def _migration_environment(database_url: str | None, extra_env: Mapping[str, str]) -> dict[str, str]: + excluded_variables: Final = ( + ("DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + if database_url is not None + else ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + ) + inherited_environment: Final = {key: value for key, value in os.environ.items() if key not in excluded_variables} + database_environment: Final = {"DATABASE_URL": database_url} if database_url is not None else {} + return { + **inherited_environment, + **database_environment, + "LITELLM_LOG": "ERROR", + **extra_env, + } + + +def _migration_invocation( + database_url: str | None, + tmp_path: Path, + extra_env: Mapping[str, str], + resolver: Literal["legacy", "v2"] = "legacy", +) -> tuple[tuple[str, ...], dict[str, str]]: + config_path: Final = tmp_path / "config.yaml" + config_path.write_text( + "model_list:\n - model_name: integration-fake\n litellm_params:\n model: openai/integration-fake\n" + ) + resolver_flag: Final = "--use_legacy_migration_resolver" if resolver == "legacy" else "--use_v2_migration_resolver" + command: Final = ( + sys.executable, + "-I", + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config_path), + resolver_flag, + "--skip_server_startup", + ) + environment: Final = _migration_environment(database_url, extra_env) + return command, environment + + +def run_v1_migrations( + database_url: str | None, + tmp_path: Path, + extra_env: Mapping[str, str], + test_name: str, + resolver: Literal["legacy", "v2"] = "legacy", +) -> MigrationResult: + command, environment = _migration_invocation(database_url, tmp_path, extra_env, resolver) + output_path: Final = _migration_log_path(test_name, tmp_path) + with output_path.open("w") as output_file: + completed: Final = subprocess.run( + command, + cwd=REPO_ROOT, + env=environment, + stdout=output_file, + stderr=subprocess.STDOUT, + text=True, + timeout=240, + ) + output: Final = output_path.read_text() + return MigrationResult(completed.returncode, output) + + +def error_lines(output: str) -> tuple[str, ...]: + return tuple(match.group(1) for match in ERROR_RECORD.finditer(output)) + + +def _json_error_records(output: str) -> tuple[dict[str, object], ...]: + records: Final = tuple(_parse_json_record(line, output) for line in output.splitlines() if line.startswith("{")) + return tuple(record for record in records if record.get("level") == "ERROR") + + +def _parse_json_record(line: str, output: str) -> dict[str, object]: + try: + record: Final = json.loads(line) + except json.JSONDecodeError: + pytest.fail(_safe_output(output)) + assert isinstance(record, dict), _safe_output(output) + return cast(dict[str, object], record) + + +def _retry_count_texts(lines: tuple[str, ...]) -> tuple[str, ...]: + return tuple(value for value in (_retry_count_text(line) for line in lines) if value is not None) + + +def _retry_count_text(line: str) -> str | None: + match: Final = RETRY_COUNT.search(line) + return match.group(1) if match is not None else None + + +def _safe_output(output: str) -> str: + return FRAGMENT_PATTERN.sub("[REDACTED]", output) + + +def _assert_no_password_fragments(output: str) -> None: + assert [fragment for fragment in FRAGMENTS if fragment in output] == [], _safe_output(output) + + +def _applied_migrations(database_url: str) -> tuple[str, ...]: + with psycopg.connect(database_url) as connection: + rows: Final = connection.execute( + 'SELECT migration_name FROM "_prisma_migrations" ' + "WHERE finished_at IS NOT NULL AND rolled_back_at IS NULL ORDER BY migration_name" + ).fetchall() + return tuple(str(row[0]) for row in rows) + + +def _duplicate_migrations(database_url: str) -> tuple[tuple[str, int], ...]: + with psycopg.connect(database_url) as connection: + rows: Final = connection.execute( + 'SELECT migration_name, COUNT(*) FROM "_prisma_migrations" ' + "GROUP BY migration_name HAVING COUNT(*) > 1 ORDER BY migration_name" + ).fetchall() + return tuple((str(row[0]), int(row[1])) for row in rows) + + +def _retry_p3018_migration(database_url: str) -> None: + agent_ids: Final = (uuid.uuid4().hex, uuid.uuid4().hex) + with psycopg.connect(database_url) as connection: + connection.execute('DROP INDEX "LiteLLM_AgentIdentity_provider_tenant_id_client_id_key"') + for agent_id in agent_ids: + connection.execute( + 'INSERT INTO "LiteLLM_AgentsTable" ' + '("agent_id", "agent_name", "agent_card_params", "created_by", "updated_by") ' + "VALUES (%s, %s, %s, %s, %s)", + (agent_id, f"audit-agent-{agent_id}", Jsonb({}), "integration", "integration"), + ) + connection.execute( + 'INSERT INTO "LiteLLM_AgentIdentity" ' + '("agent_id", "provider", "issuer", "tenant_id", "client_id", "revision") ' + "VALUES (%s, %s, %s, %s, %s, %s)", + (agent_id, "entra", f"https://audit.invalid/{agent_id}", "tenant-1", "client-1", uuid.uuid4().hex), + ) + connection.execute('DELETE FROM "_prisma_migrations" WHERE migration_name = %s', (MIGRATION_NAME,)) + + +def _relay(source: socket.socket, destination: socket.socket) -> None: + try: + while data := source.recv(65536): + destination.sendall(data) + except (BrokenPipeError, ConnectionAbortedError, ConnectionResetError): + return + + +class _PostgresForwardingServer(socketserver.ThreadingTCPServer): + allow_reuse_address = True + daemon_threads = True + request_queue_size = 64 + target: tuple[str, int] + + def __init__(self, port: int, target: tuple[str, int]) -> None: + self.target = target + super().__init__(("127.0.0.1", port), _PostgresForwardingHandler) + + +class _PostgresForwardingHandler(socketserver.BaseRequestHandler): + def handle(self) -> None: + server: Final = cast(_PostgresForwardingServer, self.server) + with socket.create_connection(server.target, timeout=10) as upstream: + reply: Final = threading.Thread(target=_relay, args=(upstream, self.request), daemon=True) + reply.start() + try: + _relay(self.request, upstream) + finally: + with suppress(OSError): + self.request.shutdown(socket.SHUT_WR) + with suppress(OSError): + upstream.shutdown(socket.SHUT_WR) + reply.join(timeout=10) + + +@contextmanager +def _gated_postgres_forwarder(port: int, target: tuple[str, int]) -> Iterator[Callable[[], None]]: + server: Final = _PostgresForwardingServer(port, target) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + try: + yield thread.start + finally: + if thread.is_alive(): + server.shutdown() + server.server_close() + if thread.ident is not None: + thread.join(timeout=10) + + +def _stop_process(process: subprocess.Popen[str]) -> None: + if process.poll() is None: + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=30) + + +def _migration_log_has_p1001_or_process_exited(output_path: Path, process: subprocess.Popen[str]) -> tuple[bool, bool]: + p1001_logged: Final = any("P1001" in line for line in error_lines(output_path.read_text())) + process_exited: Final = process.poll() is not None + return p1001_logged, process_exited + + +def test_unreachable_database_emits_four_p1001_errors_without_password_fragments( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + unreachable_url: Final = _unreachable_database_url(database_url) + completed: Final = run_v1_migrations(unreachable_url, tmp_path, {}, request.node.name) + output: Final = completed.output + errors: Final = error_lines(output) + p1001_errors: Final = tuple(line for line in errors if "P1001" in line) + assert completed.returncode == 1, _safe_output(output) + assert len(errors) == 4, _safe_output(output) + assert len(p1001_errors) == 4, _safe_output(output) + assert _retry_count_texts(errors) == (), _safe_output(output) + _assert_no_password_fragments(output) + + +def test_wrong_password_emits_four_p1000_errors_without_password_fragments( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(f"correct-{uuid.uuid4().hex}") as database_url: + parsed: Final = urlsplit(database_url) + wrong_url: Final = _database_url( + database_url, + parsed.username or "", + PASSWORD, + parsed.path.lstrip("/"), + ) + completed: Final = run_v1_migrations(wrong_url, tmp_path, {}, request.node.name) + output: Final = completed.output + errors: Final = error_lines(output) + p1000_errors: Final = tuple(line for line in errors if "P1000" in line) + assert completed.returncode == 1, _safe_output(output) + assert len(errors) == 4, _safe_output(output) + assert len(p1000_errors) == 4, _safe_output(output) + assert _retry_count_texts(errors) == (), _safe_output(output) + _assert_no_password_fragments(output) + + +def test_duplicate_agent_identity_logs_the_p3018_migration_error( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + setup: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name) + setup_output: Final = setup.output + assert setup.returncode == 0, _safe_output(setup_output) + _retry_p3018_migration(database_url) + completed: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name) + output: Final = completed.output + errors: Final = error_lines(output) + p3018_errors: Final = tuple(line for line in errors if "P3018" in line) + expected_markers: Final = ((True, True), (True, True)) + assert completed.returncode == 1, _safe_output(output) + assert len(errors) == 2, _safe_output(output) + assert tuple((MIGRATION_NAME in line, "P3018" in line) for line in p3018_errors) == expected_markers, ( + _safe_output(output) + ) + assert _retry_count_texts(errors) == ("3", "1"), _safe_output(output) + _assert_no_password_fragments(output) + + +def test_clean_database_applies_exactly_the_shipped_migrations(tmp_path: Path, request: pytest.FixtureRequest) -> None: + with owned_database(PASSWORD) as database_url: + completed: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name) + output: Final = completed.output + assert completed.returncode == 0, _safe_output(output) + assert error_lines(output) == (), _safe_output(output) + assert _applied_migrations(database_url) == SHIPPED_MIGRATIONS, _safe_output(output) + _assert_no_password_fragments(output) + + +def test_unreachable_database_recovers_after_postgres_forwarder_starts( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + admin_url: Final = os.environ["DATABASE_URL"] + admin: Final = urlsplit(admin_url) + hostname: Final = admin.hostname + port: Final = admin.port + assert hostname is not None and port is not None + target_host: Final = "127.0.0.1" if hostname == "localhost" else hostname + forwarding_port: Final = free_port() + forwarded_url: Final = _replace_port(database_url, forwarding_port, "127.0.0.1") + command, environment = _migration_invocation(forwarded_url, tmp_path, {}) + output_path: Final = _migration_log_path(request.node.name, tmp_path) + with _gated_postgres_forwarder(forwarding_port, (target_host, port)) as open_forwarder: + with output_path.open("w") as output_file: + process: Final = subprocess.Popen( + command, + cwd=REPO_ROOT, + env=environment, + stdout=output_file, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + try: + observation: Final = eventually( + lambda: _migration_log_has_p1001_or_process_exited(output_path, process), + lambda state: state[0] or state[1], + seconds=240, + ) + assert observation[0], _safe_output(output_path.read_text()) + open_forwarder() + completed_returncode: Final = process.wait(timeout=240) + finally: + _stop_process(process) + output: Final = output_path.read_text() + errors: Final = error_lines(output) + p1001_errors: Final = tuple(line for line in errors if "P1001" in line) + assert completed_returncode == 0, _safe_output(output) + assert len(p1001_errors) == 1, _safe_output(output) + assert _retry_count_texts(errors) == (), _safe_output(output) + assert _duplicate_migrations(database_url) == (), _safe_output(output) + assert _applied_migrations(database_url) == SHIPPED_MIGRATIONS, _safe_output(output) + _assert_no_password_fragments(output) + + +def test_unreachable_database_keeps_password_masked_when_shape_redaction_is_disabled( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + unreachable_url: Final = _unreachable_database_url(database_url) + completed: Final = run_v1_migrations( + unreachable_url, + tmp_path, + {"LITELLM_DISABLE_REDACT_SECRETS": "true"}, + request.node.name, + ) + output: Final = completed.output + errors: Final = error_lines(output) + p1001_errors: Final = tuple(line for line in errors if "P1001" in line) + assert completed.returncode == 1, _safe_output(output) + assert len(errors) == 4, _safe_output(output) + assert len(p1001_errors) == 4, _safe_output(output) + assert _retry_count_texts(errors) == (), _safe_output(output) + _assert_no_password_fragments(output) + + +def test_component_database_env_vars_with_wrong_password_emit_four_p1000_errors_without_password_fragments( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + correct_password: Final = f"correct-{uuid.uuid4().hex}" + with owned_database(correct_password) as database_url: + parsed: Final = urlsplit(database_url) + host: Final = parsed.hostname + port: Final = parsed.port + username: Final = parsed.username + assert host is not None and port is not None and username is not None + extra_env: Final = { + "DATABASE_HOST": f"{host}:{port}", + "DATABASE_USERNAME": username, + "DATABASE_PASSWORD": PASSWORD, + "DATABASE_NAME": parsed.path.lstrip("/"), + } + completed: Final = run_v1_migrations(None, tmp_path, extra_env, request.node.name) + output: Final = completed.output + errors: Final = error_lines(output) + p1000_errors: Final = tuple(line for line in errors if "P1000" in line) + assert completed.returncode == 1, _safe_output(output) + assert len(errors) == 4, _safe_output(output) + assert len(p1000_errors) == 4, _safe_output(output) + assert _retry_count_texts(errors) == (), _safe_output(output) + _assert_no_password_fragments(output) + + +def test_json_logs_emit_four_valid_json_p1001_error_records_without_password_fragments( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + unreachable_url: Final = _unreachable_database_url(database_url) + completed: Final = run_v1_migrations(unreachable_url, tmp_path, {"JSON_LOGS": "true"}, request.node.name) + output: Final = completed.output + errors: Final = _json_error_records(output) + messages: Final = tuple(record.get("message") for record in errors) + assert completed.returncode == 1, _safe_output(output) + assert len(errors) == 4, _safe_output(output) + assert tuple(isinstance(message, str) and "P1001" in message for message in messages) == ( + True, + True, + True, + True, + ), _safe_output(output) + assert error_lines(output) == (), _safe_output(output) + _assert_no_password_fragments(output) + + +def test_migration_job_entrypoint_emits_four_p1001_errors_without_password_fragments( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + unreachable_url: Final = _unreachable_database_url(database_url) + command: Final = (sys.executable, "-I", "-m", "litellm.proxy.prisma_migration") + environment: Final = _migration_environment( + unreachable_url, + {"USE_V2_MIGRATION_RESOLVER": "false"}, + ) + output_path: Final = _migration_log_path(request.node.name, tmp_path) + with output_path.open("w") as output_file: + completed: Final = subprocess.run( + command, + cwd=REPO_ROOT, + env=environment, + stdout=output_file, + stderr=subprocess.STDOUT, + text=True, + timeout=240, + ) + output: Final = output_path.read_text() + errors: Final = error_lines(output) + p1001_errors: Final = tuple(line for line in errors if "P1001" in line) + assert completed.returncode == 1, _safe_output(output) + assert len(errors) == 4, _safe_output(output) + assert len(p1001_errors) == 4, _safe_output(output) + _assert_no_password_fragments(output) + + +def test_v2_resolver_unreachable_database_exits_2_and_names_p1001( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + unreachable_url: Final = _unreachable_database_url(database_url) + completed: Final = run_v1_migrations(unreachable_url, tmp_path, {}, request.node.name, resolver="v2") + output: Final = completed.output + assert completed.returncode == 2, _safe_output(output) + assert "P1001" in output, _safe_output(output) + assert error_lines(output) == (), _safe_output(output) + _assert_no_password_fragments(output) + + +def test_v2_resolver_clean_database_applies_exactly_the_shipped_migrations( + tmp_path: Path, request: pytest.FixtureRequest +) -> None: + with owned_database(PASSWORD) as database_url: + completed: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name, resolver="v2") + output: Final = completed.output + assert completed.returncode == 0, _safe_output(output) + assert error_lines(output) == (), _safe_output(output) + assert _applied_migrations(database_url) == SHIPPED_MIGRATIONS, _safe_output(output) + _assert_no_password_fragments(output) diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py index 5f59fa883ce..820d033135f 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import base64 import binascii diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index 94e14798c54..5c1a996b276 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -1,6 +1,6 @@ import sys from datetime import datetime -from typing import List, Optional +from typing import Final, List, Optional import pytest from litellm._uuid import uuid import os @@ -391,6 +391,7 @@ async def test_create_mcp_server_invalid_alias(): @_SKIP_NO_MCP @pytest.mark.asyncio async def test_edit_mcp_server_redacts_credentials(): + mock_get_server: Final = mock.AsyncMock() with ( mock.patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE", @@ -399,6 +400,10 @@ async def test_edit_mcp_server_redacts_credentials(): mock.patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw" ) as mock_get_prisma, + mock.patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=mock_get_server, + ), mock.patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", new_callable=mock.AsyncMock, @@ -422,6 +427,18 @@ async def test_edit_mcp_server_redacts_credentials(): mock_manager.reload_servers_from_database = mock.AsyncMock() server_id = str(uuid.uuid4()) + stored_server: Final = LiteLLM_MCPServerTable( + server_id=server_id, + alias="Updated Server", + url="https://updated.example.com/mcp", + transport=MCPTransport.http, + created_at=datetime.now(), + updated_at=datetime.now(), + credentials={"auth_value": "secret"}, + teams=[], + ) + mock_get_server.return_value = stored_server + updated_server = LiteLLM_MCPServerTable( server_id=server_id, alias="Updated Server", @@ -458,6 +475,7 @@ async def test_edit_mcp_server_redacts_credentials(): mock_update.assert_awaited_once() mock_manager.update_server.assert_called_once_with(updated_server) mock_manager.reload_servers_from_database.assert_awaited_once() + mock_get_server.assert_awaited_once_with(mock_prisma, server_id) def test_validate_mcp_server_name_direct(): diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index a783ce6ac7e..f76aea030fe 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -1,3 +1,4 @@ +import importlib import os import socket import sys @@ -12,10 +13,10 @@ import pytest from uvicorn.importer import import_from_string from uvicorn.main import main as uvicorn_main -import gateway.main from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR, PgBouncerError, PgBouncerSettings +from litellm.proxy.proxy_server import app as proxy_app DB_ENV: Final = { "DATABASE_HOST": "db.internal", @@ -107,8 +108,22 @@ class TestUvicornArgv: argv: Final = uvicorn_argv(("--timeout-keep-alive", "30"), {"KEEPALIVE_TIMEOUT": "75"}) assert _uvicorn_params(argv)["timeout_keep_alive"] == 30 - def test_the_app_uvicorn_is_told_to_serve_is_the_trimmed_gateway(self): - assert import_from_string(cast(str, _uvicorn_params(uvicorn_argv((), {}))["app"])) is gateway.main.app + def test_the_app_uvicorn_is_told_to_serve_is_the_trimmed_gateway(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(proxy_app.router, "lifespan_context", proxy_app.router.lifespan_context) + for key in ( + "DATABASE_URL", + "DIRECT_URL", + "DATABASE_URL_READ_REPLICA", + "DATABASE_HOST", + "DATABASE_HOST_READ_REPLICA", + "DATABASE_PASSWORD", + "IAM_TOKEN_DB_AUTH", + "AZURE_POSTGRESQL_AUTH", + ): + monkeypatch.delenv(key, raising=False) + + served: Final = import_from_string(cast(str, _uvicorn_params(uvicorn_argv((), {}))["app"])) + assert served is importlib.import_module("gateway.main").app class TestPoolDatabaseUrl: diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index ea4a25283a1..70c6ec44790 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -1,11 +1,13 @@ import glob +import logging import os import re +import subprocess import sys import threading from dataclasses import dataclass from pathlib import Path -from typing import Final +from typing import Final, NoReturn, Optional import pytest @@ -19,6 +21,8 @@ sys.path.insert( from litellm_proxy_extras.utils import ( PARTITIONED_SPEND_LOGS_PUSH_ERROR, ProxyExtrasDBManager, + _redact_command_error, + _redact_credentials, filter_partitioned_spend_logs_diff, ) @@ -1412,3 +1416,411 @@ class TestMigrationJobOwnedDrift: assert 'PRIMARY KEY ("request_id")' not in filtered assert "LiteLLM_SpendLogs_legacy" not in filtered assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;' in filtered + + +_P3018_UNCLASSIFIED_STDERR: Final = ( + "Error: P3018\n\n" + "A migration failed to apply. New migrations cannot be applied before the error is " + "recovered from.\n\n" + "Migration name: 20260921190000_agent_identity\n\n" + "Database error code: 23505\n\n" + "Database error:\n" + 'ERROR: could not create unique index "agent_identity_key"\n' + "DETAIL: Key (agent_id)=(agent-1) is duplicated.\n" +) + + +_FAKE_PRISMA_PID: Final = 424242 + + +class TestV1MigrationFailuresLogAtError: + @staticmethod + def _run_v1_migrations( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + *, + deploy_stderr: Optional[str] = None, + deploy_timeout: bool = False, + diff_stderr: Optional[str] = None, + database_url: Optional[str] = None, + ) -> tuple[bool, list[list[str]], list[int]]: + import litellm_proxy_extras.utils as utils_module + + calls: Final[list[list[str]]] = [] + killed_pids: Final[list[int]] = [] + + class _FakePrismaPopen: + def __init__( + self, + argv: tuple[str, ...], + *, + env: Optional[dict[str, str]] = None, + stdout: object = None, + stderr: object = None, + text: object = None, + start_new_session: object = None, + ) -> None: + self.args: Final = argv + self.argv: Final = argv + self.pid: Final = _FAKE_PRISMA_PID + self.returncode: Optional[int] = None + calls.append(list(argv)) + + def __enter__(self) -> "_FakePrismaPopen": + return self + + def __exit__(self, *args: object) -> None: + return None + + def _subcommand(self) -> tuple[str, str]: + known: Final = ( + ("migrate", "deploy"), + ("migrate", "diff"), + ("migrate", "resolve"), + ("db", "execute"), + ) + for index in range(len(self.argv) - 1): + pair: Final = tuple(self.argv[index : index + 2]) + if pair in known: + return pair + return ("", "") + + def communicate(self, timeout: Optional[float] = None) -> tuple[str, str]: + subcommand: Final = self._subcommand() + if subcommand == ("migrate", "deploy"): + if deploy_timeout: + raise subprocess.TimeoutExpired(self.argv, timeout) + if deploy_stderr is not None: + self.returncode = 1 + return "", deploy_stderr + self.returncode = 0 + return "No pending migrations to apply", "" + if subcommand == ("migrate", "diff") and diff_stderr is not None: + self.returncode = 1 + return "", diff_stderr + self.returncode = 0 + return "", "" + + migration_dir: Final = tmp_path / "migration_dir" + migration_dir.mkdir() + if database_url is None: + monkeypatch.delenv("DATABASE_URL", raising=False) + else: + monkeypatch.setenv("DATABASE_URL", database_url) + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(migration_dir)) + monkeypatch.setattr( + utils_module.prisma_toolchain.subprocess, "Popen", _FakePrismaPopen + ) + monkeypatch.setattr( + utils_module.prisma_toolchain.os, "killpg", lambda pid, sig: killed_pids.append(pid) + ) + monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) + + succeeded: Final = ProxyExtrasDBManager._run_migrations(use_migrate=True, use_v2_resolver=False) + return succeeded, calls, killed_pids + + @staticmethod + def _deploy_call_count(calls: list[list[str]]) -> int: + return sum(1 for call in calls if tuple(call[-2:]) == ("migrate", "deploy")) + + @staticmethod + def _error_messages(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno >= logging.ERROR and record.name.startswith("litellm_proxy_extras") + ] + + def test_an_unrecognized_prisma_error_logs_its_stderr_at_error( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + stderr: Final = "Error: P1001: Can't reach database server at db:5432" + with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"): + succeeded, calls, _ = self._run_v1_migrations( + monkeypatch, tmp_path, deploy_stderr=stderr + ) + + assert succeeded is False + assert self._deploy_call_count(calls) == 4 + assert any(stderr in message for message in self._error_messages(caplog)) + + def test_an_unclassified_p3018_logs_its_stderr_and_retry_failure_at_error( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"): + succeeded, calls, _ = self._run_v1_migrations( + monkeypatch, tmp_path, deploy_stderr=_P3018_UNCLASSIFIED_STDERR + ) + + assert succeeded is False + assert self._deploy_call_count(calls) == 4 + messages: Final = self._error_messages(caplog) + assert any( + "20260921190000_agent_identity" in message and "is duplicated" in message + for message in messages + ) + assert any( + "The process failed to execute" in message and "Retrying... (3 attempts left)" in message + for message in messages + ) + + def test_called_process_error_with_no_command_retries_all_v1_attempts( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + import litellm_proxy_extras.utils as utils_module + + migration_dir: Final = tmp_path / "migration_dir" + migration_dir.mkdir() + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(migration_dir)) + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("PRISMA_OFFLINE_MODE", "true") + monkeypatch.setenv("PRISMA_CLI_PATH", sys.executable) + monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) + calls: Final[list[None]] = [] + + class _FakePrismaPopen: + def __init__( + self, + argv: tuple[str, ...], + *, + env: Optional[dict[str, str]] = None, + stdout: object = None, + stderr: object = None, + text: object = None, + start_new_session: object = None, + ) -> None: + self.args: Final = None + self.returncode: Final = 1 + calls.append(None) + + def __enter__(self) -> "_FakePrismaPopen": + return self + + def __exit__(self, *args: object) -> None: + return None + + def communicate(self, timeout: Optional[float] = None) -> tuple[str, str]: + return "", "Error: P3018 unclassified" + + monkeypatch.setattr( + utils_module.prisma_toolchain.subprocess, "Popen", _FakePrismaPopen + ) + + try: + succeeded: Final = ProxyExtrasDBManager._run_migrations( + use_migrate=True, use_v2_resolver=False + ) + except TypeError as error: + pytest.fail( + f"_run_migrations raised TypeError after {len(calls)} Popen calls: {error}", + pytrace=False, + ) + + assert succeeded is False + assert len(calls) == 4 + + def test_a_timeout_logs_at_error_naming_the_migrate_deploy_timeout_env_var( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + from litellm_proxy_extras.prisma_toolchain import PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR + + with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"): + succeeded, calls, killed_pids = self._run_v1_migrations( + monkeypatch, tmp_path, deploy_timeout=True + ) + + assert succeeded is False + assert self._deploy_call_count(calls) == 4 + assert killed_pids == [_FAKE_PRISMA_PID] * 4 + assert any( + "timed out" in message and PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR in message + for message in self._error_messages(caplog) + ) + + def test_a_recovered_baseline_logs_nothing_at_error( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"): + succeeded, calls, _ = self._run_v1_migrations( + monkeypatch, + tmp_path, + deploy_stderr=_P3005_STDERR, + database_url="postgresql://user:pass@db:5432/litellm", + ) + + assert succeeded is True + assert tuple(calls[0][-2:]) == ("migrate", "deploy") + assert self._error_messages(caplog) == [] + + def test_a_failed_baseline_recovery_logs_its_stderr_at_error( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + database_url: Final = "postgresql://llmproxy:s3cr3t 'p\"w@db:5432/litellm" + monkeypatch.delenv("DIRECT_URL", raising=False) + with caplog.at_level(logging.DEBUG, logger="litellm_proxy_extras"): + succeeded, calls, _ = self._run_v1_migrations( + monkeypatch, + tmp_path, + deploy_stderr=_P3005_STDERR, + diff_stderr=f"baseline diff failed: XYZ-7731 for {database_url}", + database_url=database_url, + ) + + assert succeeded is False + assert self._deploy_call_count(calls) == 4 + assert [ + record.getMessage() + for record in caplog.records + if "s3cr3t" in record.getMessage() or 'p"w' in record.getMessage() + ] == [] + messages: Final = self._error_messages(caplog) + assert any("postgresql://REDACTED@db:5432/litellm" in message for message in messages) + assert any("XYZ-7731" in message for message in messages) + + +@pytest.mark.parametrize( + "database_url,direct_url,text,expected", + ( + ( + "postgresql://u:pa ss@db:5432/litellm", + None, + 'Error: P1000: Authentication failed against database server at "postgresql://u:pa ss@db:5432/litellm"', + 'Error: P1000: Authentication failed against database server at "postgresql://REDACTED@db:5432/litellm"', + ), + ( + "postgresql://u:pa'ss@db:5432/litellm", + None, + "postgresql://u:pa'ss@db:5432/litellm", + "postgresql://REDACTED@db:5432/litellm", + ), + ( + 'postgresql://u:pa"ss@db:5432/litellm', + None, + 'postgresql://u:pa"ss@db:5432/litellm', + "postgresql://REDACTED@db:5432/litellm", + ), + ( + "postgresql://u:p@ss@db:5432/litellm", + None, + "postgresql://u:p@ss@db:5432/litellm", + "postgresql://REDACTED@db:5432/litellm", + ), + ( + "postgresql://u:p%20ss@db:5432/litellm", + None, + "postgresql://u:p ss@db:5432/litellm", + "postgresql://REDACTED@db:5432/litellm", + ), + ( + "postgresql://db/litellm?password=a b&sslmode=require", + None, + "postgresql://db/litellm?password=a b&sslmode=require", + "postgresql://db/litellm?REDACTED&sslmode=require", + ), + ( + "postgresql://db/litellm?sslpassword=zq'7x", + None, + "postgresql://db/litellm?sslpassword=zq'7x", + "postgresql://db/litellm?REDACTED", + ), + ( + None, + "postgresql://u:pa ss@db:5432/litellm", + "postgresql://u:pa ss@db:5432/litellm", + "postgresql://REDACTED@db:5432/litellm", + ), + ( + None, + None, + "postgresql://u:pw@db/x", + "postgresql://REDACTED@db/x", + ), + ( + "postgresql://u:p@db:5432/litellm", + None, + "Error: P1001: Can't reach database server at db:5432", + "Error: P1001: Can't reach database server at db:5432", + ), + ( + "postgresql://u:p@db:5432/litellm", + None, + "Error:P1001: Can't reach database server at db:5432", + "Error:P1001: Can't reach database server at db:5432", + ), + ( + None, + None, + "plain text with no URL", + "plain text with no URL", + ), + ), +) +def test_redact_credentials_masks_passwords_in_embedded_urls( + database_url: str | None, + direct_url: str | None, + text: str, + expected: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + if database_url is None: + monkeypatch.delenv("DATABASE_URL", raising=False) + else: + monkeypatch.setenv("DATABASE_URL", database_url) + if direct_url is None: + monkeypatch.delenv("DIRECT_URL", raising=False) + else: + monkeypatch.setenv("DIRECT_URL", direct_url) + assert _redact_credentials(text) == expected + + +@pytest.mark.parametrize("password", ("zq'7x", 'zq"7x', "zq'\"7x", "zq 7x", "zq@7x")) +def test_redact_command_error_masks_url_arguments(password: str, monkeypatch: pytest.MonkeyPatch) -> None: + database_url: Final = f"postgresql://u:{password}@db:5432/litellm" + monkeypatch.setenv("DATABASE_URL", database_url) + monkeypatch.delenv("DIRECT_URL", raising=False) + error: Final = subprocess.CalledProcessError(1, ["prisma", "migrate", "diff", "--to-url", database_url]) + + message: Final = _redact_command_error(error) + + assert "zq" not in message + assert "7x" not in message + assert "postgresql://REDACTED@db:5432/litellm" in message + assert "returned non-zero exit status 1" in message + + +@pytest.mark.parametrize( + "command", + ( + None, + Path("/usr/bin/prisma"), + 7, + ("prisma", "migrate", "deploy"), + ["prisma", "migrate", "deploy"], + "prisma migrate deploy", + ), +) +def test_redact_command_error_preserves_unredacted_command_format( + command: object, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + error: Final = subprocess.CalledProcessError(1, command) + + assert _redact_command_error(error) == str(error) + + +def test_redact_command_error_masks_password_in_tuple_url_argument( + monkeypatch: pytest.MonkeyPatch, +) -> None: + password: Final = "zq 7x" + database_url: Final = f"postgresql://u:{password}@db:5432/litellm" + monkeypatch.setenv("DATABASE_URL", database_url) + monkeypatch.delenv("DIRECT_URL", raising=False) + error: Final = subprocess.CalledProcessError(1, ("prisma", "migrate", "deploy", "--to-url", database_url)) + + message: Final = _redact_command_error(error) + + assert message.startswith("Command '('") + assert password not in message + assert "postgresql://REDACTED@db:5432/litellm" in message diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 6c8b6571991..059cff0c385 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -109,25 +109,6 @@ def set_salt_key(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") -@pytest.fixture(autouse=True) -def reset_constants_module(): - """Reset constants module to ensure clean state before each test""" - import importlib - - from litellm import constants - from litellm.proxy.auth import auth_checks - - # Reload modules before test - importlib.reload(constants) - importlib.reload(auth_checks) - - yield - - # Reload modules after test to clean up - importlib.reload(constants) - importlib.reload(auth_checks) - - @pytest.fixture def valid_sso_user_defined_values(): return LiteLLM_UserTable( @@ -875,19 +856,10 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values def test_get_cli_jwt_auth_token_custom_expiration(valid_sso_user_defined_values, monkeypatch): - """Test generating CLI JWT token with custom expiration via environment variable""" - import importlib - - from litellm import constants + """Test generating a CLI JWT token with custom expiration via the configured constant""" from litellm.proxy.auth import auth_checks - # Set custom expiration to 48 hours - monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", "48") - - # Reload the constants module to pick up the new env var - importlib.reload(constants) - # Also reload auth_checks to pick up the new constant value - importlib.reload(auth_checks) + monkeypatch.setattr(auth_checks, "CLI_JWT_EXPIRATION_HOURS", 48) token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index fc8bc289735..8b33202b483 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -7321,15 +7321,10 @@ async def test_expired_cli_session_token_is_rejected(monkeypatch): on the shared validation path, not only for DB-backed keys.""" monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test") - monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", "-1") - import importlib - - from litellm import constants from litellm.proxy.auth import auth_checks - importlib.reload(constants) - importlib.reload(auth_checks) + monkeypatch.setattr(auth_checks, "CLI_JWT_EXPIRATION_HOURS", -1) user_info = LiteLLM_UserTable( user_id="cli-admin", @@ -7346,22 +7341,17 @@ async def test_expired_cli_session_token_is_rejected(monkeypatch): mock_request.headers = {"authorization": f"Bearer {cli_token}"} mock_request.query_params = {} - try: - with ( - patch("litellm.proxy.proxy_server.master_key", "sk-master"), - patch("litellm.proxy.proxy_server.prisma_client", None), - ): - with pytest.raises(ProxyException) as exc_info: - await user_api_key_auth( - request=mock_request, - api_key=f"Bearer {cli_token}", - ) + with ( + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {cli_token}", + ) - assert exc_info.value.type == ProxyErrorTypes.expired_key - finally: - monkeypatch.delenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", raising=False) - importlib.reload(constants) - importlib.reload(auth_checks) + assert exc_info.value.type == ProxyErrorTypes.expired_key @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_custom_proxy.py b/tests/unit/proxy/test_custom_proxy.py index b646a4e80e7..a08ceccd4f3 100644 --- a/tests/unit/proxy/test_custom_proxy.py +++ b/tests/unit/proxy/test_custom_proxy.py @@ -1,52 +1,45 @@ import os +from typing import Final import uvicorn from dotenv import load_dotenv -from fastapi import FastAPI, Request +from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import JSONResponse - -load_dotenv() - -# Set the SERVER_ROOT_PATH environment variable to match the custom mount path -os.environ["SERVER_ROOT_PATH"] = "/my-custom-path" - -from litellm.proxy.proxy_server import app as litellm_app -from litellm.proxy.proxy_server import proxy_startup_event - -# Create main FastAPI app -app = FastAPI(title="Custom LiteLLM Server", lifespan=proxy_startup_event) - -# Add CORS middleware -app.add_middleware( - CORSMiddleware, - allow_origins=["*"], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -custom_path = "/my-custom-path" - -# Mount LiteLLM app at /litellm -app.mount(custom_path, litellm_app) -# Default route at / -@app.get("/") -async def root(): - return { - "message": "Welcome to the API Gateway", - "litellm_endpoint": f"{custom_path}", - } +def build_app() -> FastAPI: + load_dotenv() + os.environ["SERVER_ROOT_PATH"] = "/my-custom-path" + from litellm.proxy.proxy_server import app as litellm_app + from litellm.proxy.proxy_server import proxy_startup_event -# Health check endpoint -@app.get("/health") -async def health_check(): - return {"status": "healthy"} + app: Final = FastAPI(title="Custom LiteLLM Server", lifespan=proxy_startup_event) + custom_path: Final = "/my-custom-path" + + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + app.mount(custom_path, litellm_app) + + @app.get("/") + async def root() -> dict[str, str]: + return { + "message": "Welcome to the API Gateway", + "litellm_endpoint": custom_path, + } + + @app.get("/health") + async def health_check() -> dict[str, str]: + return {"status": "healthy"} + + return app if __name__ == "__main__": - # Run the server on port 8000 - uvicorn.run(app, host="0.0.0.0", port=4000, log_level="info") + uvicorn.run(build_app(), host="0.0.0.0", port=4000, log_level="info") diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py index 12e473f68a4..d5981b906a3 100644 --- a/tests/unit/test_constants.py +++ b/tests/unit/test_constants.py @@ -68,3 +68,36 @@ def _build_constant_env_var_map() -> dict[str, str]: env_var_map[constant_name] = env_var_name return env_var_map + + +@pytest.mark.parametrize( + ("cli_value", "litellm_cli_value", "expected"), + [ + ("48", None, 48), + (None, "48", 48), + (None, None, 24), + ("48", "72", 48), + ], + ids=("canonical-only", "alias-only", "default", "canonical-wins"), +) +def test_cli_jwt_expiration_hours_from_environment( + monkeypatch: pytest.MonkeyPatch, + cli_value: str | None, + litellm_cli_value: str | None, + expected: int, +) -> None: + monkeypatch.delenv("CLI_JWT_EXPIRATION_HOURS", raising=False) + monkeypatch.delenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", raising=False) + + try: + if cli_value is not None: + monkeypatch.setenv("CLI_JWT_EXPIRATION_HOURS", cli_value) + if litellm_cli_value is not None: + monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", litellm_cli_value) + + importlib.reload(litellm.constants) + assert litellm.constants.CLI_JWT_EXPIRATION_HOURS == expected + finally: + monkeypatch.delenv("CLI_JWT_EXPIRATION_HOURS", raising=False) + monkeypatch.delenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", raising=False) + importlib.reload(litellm.constants) diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py new file mode 100644 index 00000000000..c2aee58bb5f --- /dev/null +++ b/tests/unit/test_tool_loop.py @@ -0,0 +1,420 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.tool_loop import ToolLoopMaxRoundsExceeded +from litellm.types.llms.openai import ChatCompletionToolMessage +from litellm.types.utils import ChatCompletionMessageToolCall + +OPENAI_CHAT_COMPLETIONS_URL: Final = "https://api.openai.com/v1/chat/completions" +WEATHER_TOOLS: Final = ( + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, +) + + +def _openai_response(content: str | None, tool_calls: list | None = None) -> dict: + return { + "id": "chatcmpl-tool-loop", + "object": "chat.completion", + "created": 1739462947, + "model": "gpt-5-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls" if tool_calls else "stop", + "message": { + "role": "assistant", + "content": content, + "tool_calls": tool_calls, + }, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _tool_call(call_id: str, name: str, arguments: dict) -> dict: + return { + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": json.dumps(arguments)}, + } + + +def _tool_result(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + return ChatCompletionToolMessage(role="tool", content='{"temp": "72F"}', tool_call_id=tool_call.id or "") + + +def _request_bodies(respx_mock: respx.MockRouter) -> list[dict]: + return [json.loads(call.request.content) for call in respx_mock.calls] + + +def test_final_answer_without_tool_calls_returns_content(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=_openai_response("done")) + ) + executor_called: Final = [] + + def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executor_called.append(tc) + return _tool_result(tc) + + answer: Final = litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=executor, + api_key="sk-test", + ) + + assert answer == "done" + assert executor_called == [] + assert route.call_count == 1 + + +def test_two_rounds_appends_assistant_and_tool_messages_in_order(respx_mock: respx.MockRouter) -> None: + tool_calls: Final = [ + _tool_call("call_1", "get_weather", {"city": "Paris"}), + _tool_call("call_2", "get_weather", {"city": "Tokyo"}), + ] + route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + side_effect=[ + httpx.Response(200, json=_openai_response(None, tool_calls)), + httpx.Response(200, json=_openai_response("Paris 72F, Tokyo 60F")), + ] + ) + executed: Final = [] + + def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executed.append(tc) + return _tool_result(tc) + + messages: Final = [{"role": "user", "content": "weather in Paris and Tokyo?"}] + messages_snapshot: Final = [dict(message) for message in messages] + + answer: Final = litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=messages, + tools=WEATHER_TOOLS, + execute_tool=executor, + api_key="sk-test", + ) + + assert answer == "Paris 72F, Tokyo 60F" + assert route.call_count == 2 + assert [tc.id for tc in executed] == ["call_1", "call_2"] + assert [tc.function.name for tc in executed] == ["get_weather", "get_weather"] + assert [tc.function.arguments for tc in executed] == [ + '{"city": "Paris"}', + '{"city": "Tokyo"}', + ] + + second_body: Final = _request_bodies(respx_mock)[1] + assert second_body["messages"] == [ + {"role": "user", "content": "weather in Paris and Tokyo?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_2", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Tokyo"}'}, + }, + ], + }, + {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_1"}, + {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_2"}, + ] + + assert len(messages) == len(messages_snapshot) + assert messages == messages_snapshot + + +def test_response_format_and_tools_forwarded_every_round(respx_mock: respx.MockRouter) -> None: + respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + side_effect=[ + httpx.Response(200, json=_openai_response(None, [_tool_call("call_1", "get_weather", {"city": "Paris"})])), + httpx.Response(200, json=_openai_response('{"summary": "sunny"}')), + ] + ) + response_format: Final = { + "type": "json_schema", + "json_schema": { + "name": "weather_report", + "schema": { + "type": "object", + "properties": {"summary": {"type": "string"}}, + "required": ["summary"], + }, + }, + } + + litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "weather?"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + response_format=response_format, + api_key="sk-test", + ) + + bodies: Final = _request_bodies(respx_mock) + assert len(bodies) == 2 + for body in bodies: + assert body["response_format"] == response_format + assert body["tools"] == list(WEATHER_TOOLS) + + +def test_max_rounds_exceeded_raises_without_executing_last_round(respx_mock: respx.MockRouter) -> None: + tool_call: Final = _tool_call("call_1", "get_weather", {"city": "Paris"}) + route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=_openai_response(None, [tool_call])) + ) + executed: Final = [] + + def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executed.append(tc) + return _tool_result(tc) + + with pytest.raises(ToolLoopMaxRoundsExceeded) as exc_info: + litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "weather?"}], + tools=WEATHER_TOOLS, + execute_tool=executor, + max_rounds=2, + api_key="sk-test", + ) + + assert exc_info.value.max_rounds == 2 + assert route.call_count == 2 + assert [tc.id for tc in executed] == ["call_1"] + + +def test_max_rounds_below_one_rejected_before_any_request(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=_openai_response("done")) + ) + + with pytest.raises(ValueError, match="max_rounds must be >= 1"): + litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + max_rounds=0, + api_key="sk-test", + ) + + assert route.call_count == 0 + + +def test_stream_rejected_before_any_request(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=_openai_response("done")) + ) + + with pytest.raises(ValueError, match="stream=True is not supported"): + litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + stream=True, + api_key="sk-test", + ) + + assert route.call_count == 0 + + +def test_custom_tool_call_raises_type_error_without_executing(respx_mock: respx.MockRouter) -> None: + custom_response: Final = _openai_response( + None, + [{"id": "call_custom", "type": "custom", "custom": {"name": "apply_patch", "input": "*** patch"}}], + ) + route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=custom_response) + ) + executor_called: Final = [] + + def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executor_called.append(tc) + return _tool_result(tc) + + with pytest.raises(TypeError, match="custom tool call call_custom"): + litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=executor, + api_key="sk-test", + ) + + assert executor_called == [] + assert route.call_count == 1 + + +def _responses_payload(response_id: str, output: list) -> dict: + return { + "id": response_id, + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-5.5", + "output": output, + "parallel_tool_calls": True, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + + +def test_responses_bridge_replays_reasoning_items_across_rounds(respx_mock: respx.MockRouter) -> None: + round_one: Final = _responses_payload( + "resp_1", + [ + {"type": "reasoning", "id": "rs_abc123", "summary": [], "encrypted_content": "enc_xyz"}, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + ) + round_two: Final = _responses_payload( + "resp_2", + [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Paris is 72F", "annotations": []}], + } + ], + ) + route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[httpx.Response(200, json=round_one), httpx.Response(200, json=round_two)] + ) + executed: Final = [] + + def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executed.append(tc) + return _tool_result(tc) + + answer: Final = litellm.run_tool_loop( + model="openai/responses/gpt-5.5", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=WEATHER_TOOLS, + execute_tool=executor, + api_key="sk-test", + ) + + assert answer == "Paris is 72F" + assert route.call_count == 2 + assert [tc.id for tc in executed] == ["fc_1"] + + second_input: Final = _request_bodies(respx_mock)[1]["input"] + item_types: Final = [item.get("type") for item in second_input] + reasoning_index: Final = next(i for i, item in enumerate(second_input) if item.get("type") == "reasoning") + function_call_index: Final = next( + i for i, item in enumerate(second_input) if item.get("type") == "function_call" + ) + reasoning_item: Final = second_input[reasoning_index] + assert reasoning_item["id"] == "rs_abc123" + assert reasoning_item["encrypted_content"] == "enc_xyz" + assert reasoning_index < function_call_index, f"reasoning item must precede function_call: {item_types}" + + +async def test_arun_tool_loop_two_rounds(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler()) + tool_calls: Final = [ + _tool_call("call_1", "get_weather", {"city": "Paris"}), + _tool_call("call_2", "get_weather", {"city": "Tokyo"}), + ] + route: Final = respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + side_effect=[ + httpx.Response(200, json=_openai_response(None, tool_calls)), + httpx.Response(200, json=_openai_response("Paris 72F, Tokyo 60F")), + ] + ) + executed: Final = [] + + async def executor(tc: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executed.append(tc) + return _tool_result(tc) + + messages: Final = [{"role": "user", "content": "weather in Paris and Tokyo?"}] + messages_snapshot: Final = [dict(message) for message in messages] + + answer: Final = await litellm.arun_tool_loop( + model="openai/gpt-5-mini", + messages=messages, + tools=WEATHER_TOOLS, + execute_tool=executor, + api_key="sk-test", + ) + + assert answer == "Paris 72F, Tokyo 60F" + assert route.call_count == 2 + assert [tc.id for tc in executed] == ["call_1", "call_2"] + + second_body: Final = _request_bodies(respx_mock)[1] + assert second_body["messages"] == [ + {"role": "user", "content": "weather in Paris and Tokyo?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_2", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Tokyo"}'}, + }, + ], + }, + {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_1"}, + {"role": "tool", "content": '{"temp": "72F"}', "tool_call_id": "call_2"}, + ] + assert messages == messages_snapshot diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index b4c886752df..37027f215fe 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -42,6 +42,7 @@ "react-dom": "19.2.8", "react-error-boundary": "6.1.6", "react-hook-form": "7.82.0", + "react-hotkeys-hook": "5.3.3", "react-intersection-observer": "11.0.1", "react-json-view-lite": "2.5.0", "react-markdown": "9.1.0", @@ -10537,6 +10538,19 @@ "react": "^16.8.0 || ^17 || ^18 || ^19" } }, + "node_modules/react-hotkeys-hook": { + "version": "5.3.3", + "resolved": "https://registry.npmjs.org/react-hotkeys-hook/-/react-hotkeys-hook-5.3.3.tgz", + "integrity": "sha512-aswgyWUnE25hmhzHTfKDmKzsaSE5DJ4LKaU/o6rQSXkDd/1Bh9TfAFQbHkf6fLy11HvlYkp+cDDarGdhmCDhoQ==", + "license": "MIT", + "workspaces": [ + "packages/*" + ], + "peerDependencies": { + "react": ">=16.8.0", + "react-dom": ">=16.8.0" + } + }, "node_modules/react-intersection-observer": { "version": "11.0.1", "resolved": "https://registry.npmjs.org/react-intersection-observer/-/react-intersection-observer-11.0.1.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index bba51e9aa2b..388b6c7eab9 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -59,6 +59,7 @@ "react-dom": "19.2.8", "react-error-boundary": "6.1.6", "react-hook-form": "7.82.0", + "react-hotkeys-hook": "5.3.3", "react-intersection-observer": "11.0.1", "react-json-view-lite": "2.5.0", "react-markdown": "9.1.0", diff --git a/ui/litellm-dashboard/src/app/layout.tsx b/ui/litellm-dashboard/src/app/layout.tsx index 60db7104a03..7bb05ef11a7 100644 --- a/ui/litellm-dashboard/src/app/layout.tsx +++ b/ui/litellm-dashboard/src/app/layout.tsx @@ -7,6 +7,7 @@ import { ThemeProvider } from "next-themes"; import { AuthProvider } from "@/contexts/AuthContext"; import ReactQueryProvider from "@/contexts/ReactQueryProvider"; +import HotkeysProvider from "@/contexts/HotkeysProvider"; import { Toaster } from "@/components/ui/sonner"; const inter = Inter({ subsets: ["latin"] }); @@ -30,7 +31,9 @@ export default function RootLayout({ - {children} + + {children} + diff --git a/ui/litellm-dashboard/src/components/lens/HeaderActions.tsx b/ui/litellm-dashboard/src/components/lens/HeaderActions.tsx deleted file mode 100644 index 8df9307c990..00000000000 --- a/ui/litellm-dashboard/src/components/lens/HeaderActions.tsx +++ /dev/null @@ -1,13 +0,0 @@ -"use client"; - -import { useContext, type ReactNode } from "react"; -import { createPortal } from "react-dom"; - -import { LensPreviewContext } from "./LensPreviewButton"; - -export function HeaderActions({ children }: { children: ReactNode }) { - const preview = useContext(LensPreviewContext); - const actions =
{children}
; - if (preview === undefined) return actions; - return preview.target ? createPortal(actions, preview.target) : null; -} diff --git a/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx new file mode 100644 index 00000000000..11d02567799 --- /dev/null +++ b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx @@ -0,0 +1,101 @@ +"use client"; + +import { Tabs as TabsPrimitive } from "@base-ui/react/tabs"; +import { Activity, ScanSearch, Settings } from "lucide-react"; +import { StatusDot } from "@/components/shared/StatusDot"; +import { cn } from "@/lib/cva.config"; +import type { InvestigationActivity } from "./model/status"; +import { useWorkerConnected } from "./hooks/useWorkerConnected"; +import type { LensList } from "./model/types"; +import { LENS_TABS, type LensTab } from "./route"; +import { frameCorner, frameTab } from "./ui/frame"; + +const MODE_ICONS = { traces: Activity, investigations: ScanSearch, settings: Settings } as const; + +const ACTIVITY_DOT: Record, { className: string; label: string }> = { + running: { className: "bg-info motion-safe:animate-pulse", label: "An investigation is running" }, + queued: { className: "bg-muted-foreground/60", label: "An investigation is queued" }, +}; + +function ActivityDot({ activity }: { activity: InvestigationActivity }) { + if (activity === "idle") return null; + return ( +