diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index b617a79946c..984419717a3 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -26,6 +26,7 @@ guard_created=false guard_installed=false guard6_created=false guard6_installed=false +egress_cgroup=litellm-integration cleanup() { original_status=$? trap - EXIT INT TERM @@ -47,14 +48,14 @@ cleanup() { fi done if [ "$guard_installed" = true ]; then - sudo iptables -D OUTPUT -m owner --uid-owner "$(id -u)" -j integration_only || original_status=1 + sudo iptables -D OUTPUT -m cgroup --path "$egress_cgroup" -j integration_only || original_status=1 fi if [ "$guard_created" = true ]; then sudo iptables -F integration_only || original_status=1 sudo iptables -X integration_only || original_status=1 fi if [ "$guard6_installed" = true ]; then - sudo ip6tables -D OUTPUT -m owner --uid-owner "$(id -u)" -j integration_only || original_status=1 + sudo ip6tables -D OUTPUT -m cgroup --path "$egress_cgroup" -j integration_only || original_status=1 fi if [ "$guard6_created" = true ]; then sudo ip6tables -F integration_only || original_status=1 @@ -100,6 +101,8 @@ if [ "$mode" = parity ]; then export INTEGRATION_ROUTING=capture fi +sudo mkdir -p "/sys/fs/cgroup/$egress_cgroup" +echo "$$" | sudo tee "/sys/fs/cgroup/$egress_cgroup/cgroup.procs" > /dev/null sudo iptables -N integration_only guard_created=true sudo iptables -A integration_only -o lo -j ACCEPT @@ -109,13 +112,13 @@ for service in postgres-db redis-cache; do sudo iptables -A integration_only -d "$address" -j ACCEPT done sudo iptables -A integration_only -j REJECT -sudo iptables -I OUTPUT 1 -m owner --uid-owner "$(id -u)" -j integration_only +sudo iptables -I OUTPUT 1 -m cgroup --path "$egress_cgroup" -j integration_only guard_installed=true sudo ip6tables -N integration_only guard6_created=true sudo ip6tables -A integration_only -o lo -j ACCEPT sudo ip6tables -A integration_only -j REJECT -sudo ip6tables -I OUTPUT 1 -m owner --uid-owner "$(id -u)" -j integration_only +sudo ip6tables -I OUTPUT 1 -m cgroup --path "$egress_cgroup" -j integration_only guard6_installed=true if curl --noproxy '*' --connect-timeout 2 -s http://198.51.100.1 >/dev/null 2>&1; then diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index 37d547ff6c1..44473c9fe90 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -48,6 +48,7 @@ pub struct CyberArkSecretManager { token: Cache<(), SecretValue>, secrets: SecretCache, authentication_lock: Arc>, + policy_load_lock: Arc>, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs index 052d5570896..4b6bcfb28b0 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs @@ -26,6 +26,7 @@ impl CyberArkSecretManager { token, secrets, authentication_lock: Arc::new(tokio::sync::Mutex::new(())), + policy_load_lock: Arc::new(tokio::sync::Mutex::new(())), } } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs index 265c5fc6b28..87f6cea3619 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs @@ -1,5 +1,8 @@ use super::*; +const POLICY_LOAD_ATTEMPTS: u32 = 5; +const POLICY_LOAD_RETRY_DELAY: std::time::Duration = std::time::Duration::from_millis(200); + impl CyberArkSecretManager { pub async fn async_write_secret( &self, @@ -105,37 +108,43 @@ impl CyberArkSecretManager { "- !variable {}\n", serde_json::to_string(name).expect("serializing a string cannot fail") ); - let response = with_timeout( - self.client - .post(policy_url) - .header("Authorization", authorization) - .header("Content-Type", "application/x-yaml") - .body(body), - context, - ) - .send() - .await; - match response { - Ok(response) if response.status().is_success() => {} - Ok(response) - if matches!( - response.status(), - reqwest::StatusCode::CONFLICT | reqwest::StatusCode::UNPROCESSABLE_ENTITY - ) => - { - litellm_tracing::debug!( - "CyberArk variable policy already exists or conflicts: {}", - response.status() - ); - } - Ok(response) => { - litellm_tracing::warn!( - "Could not ensure CyberArk variable exists: {}", - response.status() - ); - } - Err(error) => { - litellm_tracing::warn!("Error ensuring CyberArk variable exists: {error}"); + let _policy_load = self.policy_load_lock.lock().await; + for attempt in 0..POLICY_LOAD_ATTEMPTS { + let response = with_timeout( + self.client + .post(policy_url.clone()) + .header("Authorization", authorization.clone()) + .header("Content-Type", "application/x-yaml") + .body(body.clone()), + context, + ) + .send() + .await; + match response { + Ok(response) + if response.status() == reqwest::StatusCode::CONFLICT + && attempt + 1 < POLICY_LOAD_ATTEMPTS => + { + tokio::time::sleep(POLICY_LOAD_RETRY_DELAY * 2_u32.pow(attempt)).await; + } + Ok(response) if response.status().is_success() => return, + Ok(response) if response.status() == reqwest::StatusCode::UNPROCESSABLE_ENTITY => { + litellm_tracing::debug!( + "CyberArk variable policy was rejected as unprocessable" + ); + return; + } + Ok(response) => { + litellm_tracing::warn!( + "Could not ensure CyberArk variable exists: {}", + response.status() + ); + return; + } + Err(error) => { + litellm_tracing::warn!("Error ensuring CyberArk variable exists: {error}"); + return; + } } } } diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs index e9a027091fa..764a591bc71 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs @@ -6,7 +6,7 @@ async fn rejected_write_token_is_reauthenticated_once() { let server = MockServer::start().await; mount_auth(&server, 2).await; Mock::given(path("/policies/acct/policy/root")) - .respond_with(ResponseTemplate::new(409)) + .respond_with(ResponseTemplate::new(201)) .expect(1) .mount(&server) .await; @@ -45,7 +45,6 @@ async fn rejected_write_token_is_reauthenticated_once() { #[rstest] #[case::created(201)] -#[case::already_exists(409)] #[case::unprocessable(422)] #[case::server_error(500)] #[tokio::test] @@ -81,13 +80,81 @@ async fn writes_tolerate_policy_status_and_cache_value(#[case] policy_status: u1 ); } +#[rstest] +#[tokio::test] +async fn policy_load_conflict_is_retried_before_the_value_write() { + let server = MockServer::start().await; + mount_auth(&server, 1).await; + let policy_loads = Arc::new(AtomicUsize::new(0)); + let policy_loads_for_response = Arc::clone(&policy_loads); + Mock::given(path("/policies/acct/policy/root")) + .respond_with(move |_: &Request| { + if policy_loads_for_response.fetch_add(1, Ordering::SeqCst) < 2 { + ResponseTemplate::new(409) + } else { + ResponseTemplate::new(201) + } + }) + .expect(3) + .mount(&server) + .await; + let policy_loads_at_value_write = Arc::clone(&policy_loads); + Mock::given(method("POST")) + .and(path("/secrets/acct/variable/key")) + .respond_with(move |_: &Request| { + if policy_loads_at_value_write.load(Ordering::SeqCst) == 3 { + ResponseTemplate::new(201) + } else { + ResponseTemplate::new(404) + } + }) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + + manager + .async_write_secret("key", &SecretValue::new("v"), None) + .await + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn concurrent_writes_load_policy_one_at_a_time() { + let server = MockServer::start().await; + mount_auth(&server, 1).await; + Mock::given(path("/policies/acct/policy/root")) + .respond_with(ResponseTemplate::new(201).set_delay(Duration::from_millis(100))) + .expect(4) + .mount(&server) + .await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(201)) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + let started = std::time::Instant::now(); + + let value = SecretValue::new("v"); + let results = tokio::join!( + manager.async_write_secret("key-0", &value, None), + manager.async_write_secret("key-1", &value, None), + manager.async_write_secret("key-2", &value, None), + manager.async_write_secret("key-3", &value, None), + ); + + assert!(results.0.is_ok() && results.1.is_ok() && results.2.is_ok() && results.3.is_ok()); + assert!(started.elapsed() >= Duration::from_millis(400)); +} + #[rstest] #[tokio::test] async fn failed_value_write_is_not_cached() { let server = MockServer::start().await; mount_auth(&server, 1).await; Mock::given(path("/policies/acct/policy/root")) - .respond_with(ResponseTemplate::new(409)) + .respond_with(ResponseTemplate::new(201)) .mount(&server) .await; Mock::given(path("/secrets/acct/variable/key")) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 2f8e7bdccea..03420b84c22 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -165,7 +165,9 @@ class LoggingWorker: if self._worker_task is None or self._worker_task.done(): self._worker_task = asyncio.create_task(self._worker_loop()) - async def _process_log_task(self, task: LoggingTask, sem: asyncio.Semaphore): + async def _process_log_task( + self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]" + ) -> None: """Runs the logging task and handles cleanup. Releases semaphore when done.""" try: if self._queue is not None: @@ -182,7 +184,7 @@ class LoggingWorker: verbose_logger.exception("LoggingWorker error: %s", e) finally: self._untrack_dequeued(task) - self._queue.task_done() + queue.task_done() finally: # Always release semaphore, even if queue is None sem.release() @@ -219,7 +221,8 @@ class LoggingWorker: async def _worker_loop(self) -> None: """Main worker loop that gets tasks and schedules them to run concurrently.""" try: - if self._queue is None or self._sem is None: + queue: Final = self._queue + if queue is None or self._sem is None: return while True: @@ -227,10 +230,10 @@ class LoggingWorker: # unbounded growth of waiting tasks await self._sem.acquire() try: - task = await self._queue.get() + task = await queue.get() self._track_dequeued(task) # Track each spawned coroutine so we can cancel on shutdown. - processing_task = asyncio.create_task(self._process_log_task(task, self._sem)) + processing_task = asyncio.create_task(self._process_log_task(task, self._sem, queue)) self._running_tasks.add(processing_task) processing_task.add_done_callback(self._running_tasks.discard) except Exception: @@ -497,7 +500,8 @@ class LoggingWorker: """ Clear the queue with a maximum time limit. """ - if self._queue is None: + queue: Final = self._queue + if queue is None: return start_time: Final = asyncio.get_event_loop().time() @@ -509,7 +513,7 @@ class LoggingWorker: break try: - task = self._queue.get_nowait() + task = queue.get_nowait() # Await the coroutine to properly execute and avoid "never awaited" warnings try: await asyncio.wait_for( @@ -522,7 +526,7 @@ class LoggingWorker: finally: # Clear reference to prevent memory leaks task = None - self._queue.task_done() # If you're using join() elsewhere + queue.task_done() except asyncio.QueueEmpty: break diff --git a/litellm/secret_managers/cyberark_secret_manager.py b/litellm/secret_managers/cyberark_secret_manager.py index b28e15c4446..f8e17488167 100644 --- a/litellm/secret_managers/cyberark_secret_manager.py +++ b/litellm/secret_managers/cyberark_secret_manager.py @@ -1,3 +1,4 @@ +import asyncio import base64 import os from typing import Any, Final @@ -10,6 +11,7 @@ import litellm from litellm._logging import verbose_logger from litellm.caching import InMemoryCache from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, _get_httpx_client, get_async_httpx_client, httpxSpecialProvider, @@ -20,6 +22,9 @@ from litellm.rust_bridge.secret_manager import resolve_native_provider_reader, r from .base_secret_manager import BaseSecretManager, raise_if_unsafe_secret_name from .main import str_to_bool +CYBERARK_POLICY_LOAD_ATTEMPTS: Final = 5 +CYBERARK_POLICY_LOAD_RETRY_DELAY_SECONDS: Final = 0.2 + class CyberArkSecretManager(BaseSecretManager): def __init__(self): @@ -30,6 +35,7 @@ class CyberArkSecretManager(BaseSecretManager): self.conjur_account = os.getenv("CYBERARK_ACCOUNT", "default") self.conjur_username = os.getenv("CYBERARK_USERNAME", "admin") self.conjur_api_key = os.getenv("CYBERARK_API_KEY", "") + self._policy_load_lock: Final = asyncio.Lock() # Optional config for certificate-based auth self.tls_cert_path = os.getenv("CYBERARK_CLIENT_CERT", "") @@ -118,7 +124,7 @@ class CyberArkSecretManager(BaseSecretManager): token: Final = self._authenticate() return {"Authorization": f'Token token="{token}"'} - def _ensure_variable_exists(self, secret_name: str) -> None: + async def _ensure_variable_exists(self, secret_name: str, async_client: AsyncHTTPHandler) -> None: """ Ensure a variable exists in CyberArk Conjur by creating a policy entry if needed. @@ -134,27 +140,33 @@ class CyberArkSecretManager(BaseSecretManager): policy_yaml: Final = f"- !variable {quoted_name}\n" try: - client: Final = _get_httpx_client(params={"ssl_verify": self.ssl_verify}) - resp: Final = client.client.post( - policy_url, - headers={ - **self._get_request_headers(), - "Content-Type": "application/x-yaml", - }, - content=policy_yaml, - ) - resp.raise_for_status() - verbose_logger.debug("Created policy entry for variable: %s", secret_name) - except httpx.HTTPStatusError as e: - # Variable might already exist, which is fine - if e.response.status_code in [409, 422]: - verbose_logger.debug("Variable %s already exists or policy conflict (expected)", secret_name) - else: - verbose_logger.warning( - "Could not ensure variable exists: %s - %s", e.response.status_code, e.response.text - ) + async with self._policy_load_lock: + resp: Final = await self._load_variable_policy(async_client, policy_url, policy_yaml) except Exception as e: verbose_logger.warning("Error ensuring variable exists: %s", e) + return + if resp.is_success: + verbose_logger.debug("Created policy entry for variable: %s", secret_name) + elif resp.status_code == 422: + verbose_logger.debug("Variable %s policy was rejected as unprocessable", secret_name) + else: + verbose_logger.warning("Could not ensure variable exists: %s - %s", resp.status_code, resp.text) + + async def _load_variable_policy( + self, async_client: AsyncHTTPHandler, policy_url: str, policy_yaml: str, attempt: int = 0 + ) -> httpx.Response: + resp: Final = await async_client.client.post( + policy_url, + headers={ + **self._get_request_headers(), + "Content-Type": "application/x-yaml", + }, + content=policy_yaml, + ) + if resp.status_code != 409 or attempt + 1 == CYBERARK_POLICY_LOAD_ATTEMPTS: + return resp + await asyncio.sleep(CYBERARK_POLICY_LOAD_RETRY_DELAY_SECONDS * (1 << attempt)) + return await self._load_variable_policy(async_client, policy_url, policy_yaml, attempt + 1) def get_url(self, secret_name: str) -> str: """ @@ -303,7 +315,7 @@ class CyberArkSecretManager(BaseSecretManager): try: # Ensure the variable exists in the policy first - self._ensure_variable_exists(secret_name) + await self._ensure_variable_exists(secret_name, async_client) # Now set the secret value url: Final = self.get_url(secret_name) diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index 86e47c0b1e1..f889844f1ae 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -121,7 +121,8 @@ def cleanup_batch( if current.status == "cancelling" and not needs_terminal_state: return assert clock() < deadline, ( - f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" + f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, " + f"last status {current.status}" ) wait(BATCH_CANCEL_POLL_SECONDS) diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index a0932a80dfe..28eb362e876 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -210,6 +210,7 @@ class TestBatchCancellation: with pytest.raises(ExceptionGroup) as caught: manager.teardown() assert "cancellation did not finish" in str(caught.value.exceptions[0]) + assert "last status cancelling" in str(caught.value.exceptions[0]) client.calls.assert_done() @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py index 67fd88bc84d..c3697a31424 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py @@ -74,6 +74,8 @@ SPEND_ROUTES = ( _SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity") +_CAPTURE_RATE_ROUTE: Final = "/spend/capture_rate" + # Served from the MonthlyGlobalSpend / DailyTagSpend / Last30d* views, which the # proxy creates in the background once the schema migrations have landed, so on a # fresh database they can 500 for a while after the proxy starts serving. @@ -120,7 +122,7 @@ def test_schema_listed_spend_routes_are_responsive(client: SpendClient) -> None: and "{" not in path and any(path.startswith(prefix) for prefix in _SPEND_PREFIXES) ] - extras = [path for path in discovered if path not in SPEND_ROUTES] + extras = [path for path in discovered if path not in (*SPEND_ROUTES, _CAPTURE_RATE_ROUTE)] params = _date_range() results = [(path, client.probe(path, params=params)) for path in extras] @@ -132,3 +134,11 @@ def test_schema_listed_spend_routes_are_responsive(client: SpendClient) -> None: if not result.healthy ] assert not offenders, "non-responsive schema spend routes:\n" + "\n".join(offenders) + + +def test_capture_rate_reports_or_names_the_missing_billing_key(client: SpendClient) -> None: + result: Final = client.probe(_CAPTURE_RATE_ROUTE, params=_date_range()) + print(f"{_CAPTURE_RATE_ROUTE} -> {result.status_code}\n{result.body[:600]}") + assert result.status_code == 200 or (result.status_code == 503 and "OPENAI_ADMIN_KEY is not set" in result.body), ( + f"{_CAPTURE_RATE_ROUTE} -> {result.status_code}\n{result.body[:600]}" + ) diff --git a/tests/e2e/quota_management/spend_tracking/test_websearch_interception_session_e2e.py b/tests/e2e/quota_management/spend_tracking/test_websearch_interception_session_e2e.py index 89d0beec414..6352ab67c3c 100644 --- a/tests/e2e/quota_management/spend_tracking/test_websearch_interception_session_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_websearch_interception_session_e2e.py @@ -12,13 +12,14 @@ Needs a proxy booted with the callback and a real search backend, which ``gateway/stage_mirror_ci_config.yml`` carries as the ``e2e-search`` Perplexity tool. """ -from typing import Final +from typing import Final, Literal import pytest from e2e_config import unique_marker from e2e_http import unwrap from lifecycle import ResourceManager from models import ( + AnthropicContentBlock, AnthropicMessagesBody, AnthropicWebSearchTool, ChatMessage, @@ -26,6 +27,7 @@ from models import ( SpendLogRow, ) from proxy_client import ProxyClient +from pydantic import BaseModel, ValidationError pytestmark = pytest.mark.e2e @@ -37,6 +39,18 @@ def _has_search_row(rows: list[SpendLogRow]) -> bool: return any(row.call_type == SEARCH_CALL_TYPE for row in rows) +class _SearchResultError(BaseModel): + type: Literal["web_search_tool_result_error"] + error_code: str + + +def _search_error_code(block: AnthropicContentBlock) -> str | None: + try: + return _SearchResultError.model_validate((block.model_extra or {}).get("content")).error_code + except ValidationError: + return None + + class TestWebSearchInterceptionSession: @pytest.mark.covers( "quota_management.spend_tracking.websearch_interception.bills_under_request_session", @@ -79,6 +93,15 @@ class TestWebSearchInterceptionSession: f"precondition: the turn never ran an intercepted search, so there is no search row to attribute. " f"blocks={block_types}" ) + search_errors: Final = tuple( + code + for block in response.content or () + if block.type == "web_search_tool_result" and (code := _search_error_code(block)) is not None + ) + assert not search_errors, ( + f"precondition: the e2e-search tool failed upstream ({search_errors}), so no {SEARCH_CALL_TYPE} row is " + "billed at all; check the proxy's search tool credentials before reading this as a session bug" + ) rows: Final = proxy.poll_logs_for_session(session_id, min_rows=2, predicate=_has_search_row) by_call_type: Final = {row.call_type or "" for row in rows} diff --git a/tests/e2e/secret_manager/secret_store_cyberark.py b/tests/e2e/secret_manager/secret_store_cyberark.py index 87bcd2ffb1c..375b997e02c 100644 --- a/tests/e2e/secret_manager/secret_store_cyberark.py +++ b/tests/e2e/secret_manager/secret_store_cyberark.py @@ -2,6 +2,7 @@ from __future__ import annotations import base64 import os +import time from dataclasses import dataclass, field from typing import Final, Literal from urllib.parse import quote @@ -25,6 +26,9 @@ DEFAULT_USERNAME: Final = "admin" SYSTEM: Final = "cyberark" +_POLICY_LOAD_ATTEMPTS: Final = 5 +_POLICY_LOAD_RETRY_DELAY_SECONDS: Final = 0.2 + _START_HINT: Final = ( f"Start one with `bash tests/e2e/secret_manager/backend.sh up {SYSTEM}`, which writes the env for " f"the proxy (booted from gateway/secret_manager_{SYSTEM}_ci_config.yml) and for the tests" @@ -72,13 +76,20 @@ class Conjur: def _secret_url(self, name: str) -> str: return f"{self.base_url}/secrets/{self.account}/variable/{quote(name, safe='')}" - def _update_root_policy(self, method: Literal["POST", "PATCH"], policy: str, action: str) -> None: + def _load_root_policy(self, method: Literal["POST", "PATCH"], policy: str, attempt: int = 0) -> ExternalWrite: result: Final = send_text_external( method, f"{self.base_url}/policies/{self.account}/policy/root", headers=self._headers(content_type="application/x-yaml"), content=policy, ) + if result.status_code != 409 or attempt + 1 == _POLICY_LOAD_ATTEMPTS: + return result + time.sleep(_POLICY_LOAD_RETRY_DELAY_SECONDS * (1 << attempt)) + return self._load_root_policy(method, policy, attempt + 1) + + def _update_root_policy(self, method: Literal["POST", "PATCH"], policy: str, action: str) -> None: + result: Final = self._load_root_policy(method, policy) self._fail_unless_reached(result, action) if not result.ok: pytest.fail(f"Conjur refused to {action}: HTTP {result.status_code} {result.body[:300]}") diff --git a/tests/e2e/ui/tests/auth/logout.spec.ts b/tests/e2e/ui/tests/auth/logout.spec.ts index 92c31456353..fcdd71898e2 100644 --- a/tests/e2e/ui/tests/auth/logout.spec.ts +++ b/tests/e2e/ui/tests/auth/logout.spec.ts @@ -24,6 +24,11 @@ test.describe("Logout", () => { // Click Logout — the handler clears the auth cookie and navigates via // window.location.href = PROXY_LOGOUT_URL (empty string in the e2e env). await popup.getByRole("button", { name: "Logout" }).click(); + await expect + .poll(async () => (await page.context().cookies()).filter((c) => c.name === "token").length, { + timeout: 15_000, + }) + .toBe(0); // The cookie is now gone — visiting a protected page must redirect to /ui/login. await page.goto("/ui?page=llm-playground", { waitUntil: "domcontentloaded" }); diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 6beea9ae8f4..26f5ced4de7 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -47,6 +47,8 @@ def _model_double(tool: str) -> Callable[[Request], Reply]: arguments: Final = json.dumps(ADD) def respond(request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return _json({"object": "list", "data": []}) body: Final = json.loads(request.body) assert isinstance(body, dict), request.body done: Final = _has_tool_result(body) @@ -192,7 +194,9 @@ class Rig: ) def upstream_tools(self) -> tuple[tuple[str, ...], ...]: - return tuple(_tool_names(json.loads(request.body)) for request in self.wire.drain()) + return tuple( + _tool_names(json.loads(request.body)) for request in self.wire.drain() if request.method == "POST" + ) def final_text(self, body: Mapping[str, object]) -> str: if self.surface == "chat": diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py index d94b3b24954..611bd6a1264 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_chaos.py +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -2,6 +2,7 @@ import asyncio import json import re import signal +from dataclasses import dataclass from pathlib import Path from typing import Final @@ -51,6 +52,16 @@ def _error_information(call_id: str) -> dict[str, JsonValue]: return object_value(parsed["error_information"]) +@dataclass(frozen=True, slots=True) +class _Served: + response: httpx.Response + client_port: int + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)) + + def _single_spend_row(call_id: str) -> None: rows: Final = eventually( lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), @@ -62,19 +73,23 @@ def _single_spend_row(call_id: str) -> None: async def _fire_burst( base_url: str, key: str, count: int, *, tolerate_transport_errors: bool = False -) -> tuple[httpx.Response, ...]: - async def one(client: httpx.AsyncClient, index: int) -> httpx.Response: +) -> tuple[_Served, ...]: + async def one(client: httpx.AsyncClient, index: int) -> _Served: if index % 3 == 0: path: Final = "/gemini/v1beta/models/nope-9:generateContent" elif index % 3 == 1: path = "/gemini/v1beta/models/nope-9:streamGenerateContent?alt=sse" else: path = "/gemini/v1beta/models/healthy-model:streamGenerateContent?alt=sse" - return await client.post( + async with client.stream( + "POST", path, json=_GENERATE_CONTENT, headers={"Authorization": f"Bearer {key}", "x-goog-api-key": key}, - ) + ) as response: + client_port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) + await response.aread() + return _Served(response=response, client_port=client_port) async with httpx.AsyncClient(base_url=base_url, timeout=30, trust_env=False) as client: results: Final = await asyncio.gather( @@ -82,7 +97,7 @@ async def _fire_burst( ) for result in results: assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) - return tuple(result for result in results if isinstance(result, httpx.Response)) + return tuple(result for result in results if isinstance(result, _Served)) async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gateway: Gateway, tmp_path: Path) -> None: @@ -97,7 +112,7 @@ async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gate burst: Final = asyncio.create_task(_fire_burst(str(candidate.client.base_url), candidate.key, 30)) await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 10, 30) with wire_server(_chaos_reply, port=port): - responses: Final = await burst + responses: Final = tuple(served.response for served in await burst) assert len(responses) == 30 for response in responses: assert response.status_code in (200, 404, 500, 502), response.status_code @@ -131,10 +146,15 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat _fire_burst(str(candidate.client.base_url), candidate.key, 20, tolerate_transport_errors=True) ) await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 5, 30) - psutil.Process(workers[0]).send_signal(signal.SIGKILL) - responses: Final = await burst - for response in responses: - assert response.status_code in (200, 404, 500, 502), response.status_code + victim: Final = psutil.Process(workers[0]) + victim.suspend() + victim_ports: Final = frozenset( + connection.raddr.port for connection in victim.net_connections(kind="tcp") if connection.raddr + ) + victim.send_signal(signal.SIGKILL) + served: Final = await burst + for item in served: + assert item.response.status_code in (200, 404, 500, 502), item.response.status_code follow_up: Final = candidate.request( "POST", "/gemini/v1beta/models/nope-9:generateContent", @@ -143,8 +163,12 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat ) assert follow_up.status_code == 404, follow_up.text assert follow_up.json() == json.loads(_NOT_FOUND_BODY), follow_up.text - for response in responses: - if "x-litellm-call-id" in response.headers: - _single_spend_row(response.headers["x-litellm-call-id"]) + logged: Final = tuple(item for item in served if "x-litellm-call-id" in item.response.headers) + survivor_served: Final = tuple(item for item in logged if item.client_port not in victim_ports) + assert survivor_served, [item.client_port for item in logged] + for item in survivor_served: + _single_spend_row(item.response.headers["x-litellm-call-id"]) + for item in logged: + assert len(_spend_rows(item.response.headers["x-litellm-call-id"])) <= 1, item.response.headers error_information: Final = _error_information(follow_up.headers["x-litellm-call-id"]) assert "not found for this scripted upstream" in str(error_information["error_message"]), follow_up.text diff --git a/tests/litellm_utils_tests/test_cyberark.py b/tests/litellm_utils_tests/test_cyberark.py index 9172e33af10..6d52cd9b079 100644 --- a/tests/litellm_utils_tests/test_cyberark.py +++ b/tests/litellm_utils_tests/test_cyberark.py @@ -86,7 +86,8 @@ async def test_cyberark_write_secret_rejects_yaml_injection(): "team/user@example.com", ], ) -def test_cyberark_ensure_variable_exists_escapes_yaml_metacharacters(secret_name): +@pytest.mark.asyncio +async def test_cyberark_ensure_variable_exists_escapes_yaml_metacharacters(secret_name): """ Regression test: _ensure_variable_exists must escape secret_name (not just denylist-check it) so the policy body always parses back to exactly one @@ -95,19 +96,21 @@ def test_cyberark_ensure_variable_exists_escapes_yaml_metacharacters(secret_name with patch("litellm.proxy.proxy_server.premium_user", True): captured = {} - def _capture_post(url, headers=None, content=None): + async def _capture_post(url, headers=None, content=None): captured["content"] = content return create_mock_response(status_code=201, text="") mock_sync_client = MagicMock() - mock_sync_client.client.post.side_effect = _capture_post + mock_sync_client.client.post.return_value = create_mock_response(status_code=200, text="mock-token") + mock_async_client = MagicMock() + mock_async_client.client.post.side_effect = _capture_post with patch( "litellm.secret_managers.cyberark_secret_manager._get_httpx_client", return_value=mock_sync_client, ): cyberark_manager = CyberArkSecretManager() - cyberark_manager._ensure_variable_exists(secret_name) + await cyberark_manager._ensure_variable_exists(secret_name, mock_async_client) policy_yaml = captured["content"] parsed = yaml.compose(policy_yaml) diff --git a/tests/local_testing/test_alangfuse.py b/tests/local_testing/test_alangfuse.py index a9d111843fd..ec80724d3ba 100644 --- a/tests/local_testing/test_alangfuse.py +++ b/tests/local_testing/test_alangfuse.py @@ -5,6 +5,10 @@ import logging import os from typing import Any, Optional from unittest.mock import MagicMock, patch +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer + +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest logging.basicConfig(level=logging.DEBUG) @@ -206,53 +210,91 @@ def create_async_task(**completion_kwargs): return asyncio.create_task(litellm.acompletion(**completion_args)) +def _otlp_capture(exports: list[bytes]) -> type[BaseHTTPRequestHandler]: + class OtlpCapture(BaseHTTPRequestHandler): + def do_POST(self): + exports.append(self.rfile.read(int(self.headers.get("content-length", 0)))) + self.send_response(200) + self.end_headers() + + def do_GET(self): + self.send_response(200) + self.send_header("content-type", "application/json") + self.end_headers() + self.wfile.write(b"{}") + + def log_message(self, *args): + pass + + return OtlpCapture + + +@pytest.fixture +def local_langfuse(): + exports: list[bytes] = [] + server = HTTPServer(("127.0.0.1", 0), _otlp_capture(exports)) + threading.Thread(target=server.serve_forever, daemon=True).start() + yield f"http://127.0.0.1:{server.server_port}", exports + server.shutdown() + + +def _exported_spans(exports: list[bytes]): + for body in exports: + for resource_spans in ExportTraceServiceRequest.FromString(body).resource_spans: + for scope_spans in resource_spans.scope_spans: + yield from scope_spans.spans + + +def _exported_attributes(exports: list[bytes], trace_id: str) -> list[dict[str, str]]: + return [ + {attribute.key: attribute.value.string_value for attribute in span.attributes} + for span in _exported_spans(list(exports)) + if span.trace_id.hex() == trace_id + ] + + @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.flaky(retries=12, delay=2) -async def test_langfuse_logging_without_request_response(stream, langfuse_client): - try: - from litellm._uuid import uuid +async def test_langfuse_logging_without_request_response(stream, local_langfuse, monkeypatch): + from litellm._uuid import uuid - _unique_trace_name = f"litellm-test-{str(uuid.uuid4())}" - litellm.set_verbose = True - litellm.turn_off_message_logging = True - litellm.success_callback = ["langfuse"] - response = await create_async_task( - model="gpt-3.5-turbo", - stream=stream, - metadata={"trace_id": _unique_trace_name}, - ) - print(response) - if stream: - async for chunk in response: - print(chunk) + langfuse_host, exports = local_langfuse + prompt = f"prompt-{uuid.uuid4()}" + answer = f"answer-{uuid.uuid4()}" + trace_name = f"litellm-test-{uuid.uuid4()}" + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + monkeypatch.setattr(litellm, "success_callback", ["langfuse"]) + response = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": prompt}], + mock_response=answer, + stream=stream, + metadata={"trace_id": trace_name}, + langfuse_public_key=f"pk-lf-{trace_name}", + langfuse_secret_key="sk-lf-local", + langfuse_host=langfuse_host, + ) + if stream: + async for _ in response: + pass - langfuse_client.flush() + generations: list[dict[str, str]] = [] + for _ in range(60): + generations = [ + attributes + for attributes in _exported_attributes(exports, resolve_trace_id(trace_name)) + if attributes.get("langfuse.observation.type") == "generation" + ] + if generations: + break + await asyncio.sleep(0.5) - for _ in range(30): - _trace_data = langfuse_client.api.observations.get_many( - trace_id=resolve_trace_id(_unique_trace_name), - type="GENERATION", - fields="core,io", - ).data - if _trace_data: - break - await asyncio.sleep(3) - - print(f"_trace_data: {_trace_data}") - assert json.loads(_trace_data[0].input) == { - "messages": [{"content": "redacted-by-litellm", "role": "user"}] - } - assert json.loads(_trace_data[0].output) == { - "role": "assistant", - "content": "redacted-by-litellm", - "function_call": None, - "tool_calls": None, - "provider_specific_fields": None, - } - - except Exception as e: - pytest.fail(f"An exception occurred - {e}") + assert len(generations) == 1, generations + assert json.loads(generations[0]["langfuse.observation.input"]) == { + "messages": [{"content": "redacted-by-litellm", "role": "user"}] + } + assert json.loads(generations[0]["langfuse.observation.output"])["content"] == "redacted-by-litellm" + assert all(prompt.encode() not in body and answer.encode() not in body for body in exports) # Get the current directory of the file being run diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 7e593767ea9..d3ad1d989c8 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -4,7 +4,7 @@ import os import traceback from dotenv import load_dotenv from fastapi import Request -from datetime import datetime +from datetime import datetime, timezone from litellm import Router import pytest @@ -971,11 +971,18 @@ def _rpm_tpm_router(model_id: str) -> Router: ) +@pytest.fixture +def router_minute_pinned(monkeypatch): + pinned = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) + monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned) + + def _ratelimit_headers(response: ModelResponse | CustomStreamWrapper) -> dict[str, int]: return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")} @pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") async def test_acompletion_headers_read_post_increment_counter_and_count_once(): router = _rpm_tpm_router("lit-3058-async") @@ -1018,6 +1025,7 @@ async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deplo @pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion(): router = _rpm_tpm_router("lit-3058-stream") diff --git a/tests/unit/a2a_protocol/test_cost_calculator.py b/tests/unit/a2a_protocol/test_cost_calculator.py index 56d3d57c89e..8d8ec815f3a 100644 --- a/tests/unit/a2a_protocol/test_cost_calculator.py +++ b/tests/unit/a2a_protocol/test_cost_calculator.py @@ -10,6 +10,12 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + +async def _reset_callbacks_and_settle_pending_logs() -> None: + litellm.logging_callback_manager._reset_all_callbacks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) def _make_send_message_request(request_id: str, user_text: str = "Hello"): @@ -129,7 +135,7 @@ async def test_asend_message_uses_cost_per_query(monkeypatch): from litellm.a2a_protocol import asend_message # Setup logger - litellm.logging_callback_manager._reset_all_callbacks() + await _reset_callbacks_and_settle_pending_logs() cost_logger = CostLogger() monkeypatch.setattr(litellm, "callbacks", [cost_logger]) @@ -164,7 +170,7 @@ async def test_asend_message_uses_cost_per_query_from_litellm_params_dict(monkey """ from litellm.a2a_protocol import asend_message - litellm.logging_callback_manager._reset_all_callbacks() + await _reset_callbacks_and_settle_pending_logs() cost_logger = CostLogger() monkeypatch.setattr(litellm, "callbacks", [cost_logger]) @@ -225,7 +231,7 @@ async def test_asend_message_uses_input_output_cost_per_token(monkeypatch): from litellm.a2a_protocol import asend_message # Setup logger - litellm.logging_callback_manager._reset_all_callbacks() + await _reset_callbacks_and_settle_pending_logs() token_cost_logger = TokenAndCostLogger() monkeypatch.setattr(litellm, "callbacks", [token_cost_logger]) @@ -299,7 +305,7 @@ async def test_asend_message_passes_agent_id_to_callback(monkeypatch): from litellm.a2a_protocol import asend_message # Setup logger - litellm.logging_callback_manager._reset_all_callbacks() + await _reset_callbacks_and_settle_pending_logs() agent_id_logger = AgentIdLogger() monkeypatch.setattr(litellm, "callbacks", [agent_id_logger]) @@ -359,7 +365,7 @@ async def test_asend_message_streaming_propagates_metadata(): from litellm.a2a_protocol import asend_message_streaming # Setup logger - litellm.logging_callback_manager._reset_all_callbacks() + await _reset_callbacks_and_settle_pending_logs() metadata_logger = MetadataLogger() litellm.logging_callback_manager.add_litellm_async_success_callback(metadata_logger) @@ -406,7 +412,7 @@ async def test_asend_message_streaming_triggers_callbacks(): from litellm.a2a_protocol import asend_message_streaming # Setup logger - must use logging_callback_manager to properly register - litellm.logging_callback_manager._reset_all_callbacks() + await _reset_callbacks_and_settle_pending_logs() callback_logger = AgentIdLogger() litellm.logging_callback_manager.add_litellm_async_success_callback(callback_logger) litellm.logging_callback_manager.add_litellm_success_callback(callback_logger) diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 5d72fe7213d..5f83be7c7bc 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -2,6 +2,7 @@ import asyncio import time from collections.abc import Iterator from datetime import timedelta +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1023,7 +1024,12 @@ async def test_breaker_metrics_track_state_and_failure_class(): from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure + from litellm.caching.redis_cache import RedisCircuitBreaker, _breaker_metrics, is_redis_timeout_failure + + metrics: Final = _breaker_metrics() + for collector in (metrics._state_gauge, metrics._transitions, metrics._failures): + if collector is not None and collector not in REGISTRY._collector_to_names: + REGISTRY.register(collector) def sample(name, labels=None): return REGISTRY.get_sample_value(name, labels) or 0.0 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index ecea4723bf4..ec957d80904 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -12,6 +12,32 @@ import httpx import pytest from pytest_socket import enable_socket, socket_allow_hosts +HOST_ENVIRONMENT_ALLOWLIST: Final = frozenset( + ( + "PATH", + "HOME", + "USER", + "LOGNAME", + "TMPDIR", + "TEMP", + "TMP", + "LANG", + "LC_ALL", + "LC_CTYPE", + "TZ", + "VIRTUAL_ENV", + "LITELLM_LOCAL_MODEL_COST_MAP", + "TIKTOKEN_CACHE_DIR", + ) +) +HOST_ENVIRONMENT_ALLOWED_PREFIXES: Final = ("PYTEST_", "PYTHON", "COV_CORE_", "COVERAGE_") +HOST_ONLY_ENVIRONMENT: Final = frozenset( + name + for name in os.environ + if name not in HOST_ENVIRONMENT_ALLOWLIST and not name.startswith(HOST_ENVIRONMENT_ALLOWED_PREFIXES) +) + +os.environ["PYTHON_DOTENV_DISABLED"] = "1" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import @@ -170,6 +196,8 @@ def isolated_aws_config_files(tmp_path_factory: pytest.TempPathFactory) -> tuple def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> Iterator[None]: credentials, config = isolated_aws_config_files with pytest.MonkeyPatch.context() as environment: + for name in HOST_ONLY_ENVIRONMENT: + environment.delenv(name, raising=False) environment.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials)) environment.setenv("AWS_CONFIG_FILE", str(config)) environment.setenv("AWS_EC2_METADATA_DISABLED", "true") diff --git a/tests/unit/enterprise/integrations/test_prometheus.py b/tests/unit/enterprise/integrations/test_prometheus.py index 7315f2b9881..16d6ff9d9a0 100644 --- a/tests/unit/enterprise/integrations/test_prometheus.py +++ b/tests/unit/enterprise/integrations/test_prometheus.py @@ -477,7 +477,7 @@ def test_valid_configuration_passes_validation(): # ============================================================================== -@pytest.fixture +@pytest.fixture(autouse=True) def reset_prometheus_exclude_settings(): """Restore the global exclude settings after each test so they don't leak.""" prev_metrics = litellm.prometheus_exclude_metrics diff --git a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py index 7ef43e2eadf..21e50561f57 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py @@ -5,7 +5,6 @@ Tests the end-to-end flow of websearch_interception callback with litellm.acompletion() for transparent server-side web search execution. """ -import os from unittest.mock import MagicMock import pytest @@ -37,75 +36,6 @@ def websearch_logger(): return WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI, LlmProviders.MINIMAX]) -@pytest.mark.asyncio -@pytest.mark.skipif( - os.environ.get("OPENAI_API_KEY") is None, - reason="OPENAI_API_KEY not set", -) -async def test_websearch_chat_completion_with_openai(): - """Test websearch interception with OpenAI chat completions API. - - This test verifies that: - 1. Model calls litellm_web_search tool - 2. Server executes web search automatically - 3. Server makes follow-up request with search results - 4. User gets final answer without tool_calls - """ - # Configure WebSearch interception - original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] - websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) - litellm.callbacks = [websearch_logger] - - try: - response = await litellm.acompletion( - model="gpt-4o-mini", # Use cheaper model for testing - messages=[ - { - "role": "user", - "content": "What's the weather in San Francisco today?", - } - ], - tools=[ - { - "type": "function", - "function": { - "name": "litellm_web_search", - "description": "Search the web for information", - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query", - } - }, - "required": ["query"], - }, - }, - } - ], - ) - - # Verify response structure - assert isinstance(response, ModelResponse) - assert response.choices[0].message.content is not None - assert len(response.choices[0].message.content) > 0 - - # If agentic loop worked, we should NOT have tool_calls in final response - # (they should have been executed and replaced with final answer) - if hasattr(response.choices[0].message, "tool_calls"): - # If tool_calls exist, it means agentic loop didn't run - # This could happen if search tool is not configured - pytest.skip("Agentic loop did not execute - search tool may not be configured") - - # Verify we got a meaningful response - assert response.choices[0].finish_reason in ["stop", "end_turn"] - - finally: - # Restore original callbacks - litellm.callbacks = original_callbacks - - @pytest.mark.asyncio async def test_websearch_chat_completion_hook_detection(): """Test that websearch hook correctly detects tool calls in response.""" @@ -321,61 +251,6 @@ async def test_websearch_json_serialization_fix(): assert arguments_str != "{'query': 'weather in SF'}" -@pytest.mark.asyncio -@pytest.mark.skipif( - os.environ.get("OPENAI_API_KEY") is None or os.environ.get("PERPLEXITY_API_KEY") is None, - reason="OPENAI_API_KEY or PERPLEXITY_API_KEY not set", -) -async def test_websearch_streaming_conversion(): - """Test that streaming requests are converted to non-streaming for web search. - - When stream=True is passed with web search tools, the handler should: - 1. Convert stream=True to stream=False for initial request - 2. Execute web search - 3. Convert final response back to streaming - """ - websearch_logger = WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.OPENAI], search_tool_name="perplexity-search" - ) - litellm.callbacks = [websearch_logger] - - try: - response = await litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "What's the latest AI news?"}], - tools=[ - { - "type": "function", - "function": { - "name": "litellm_web_search", - "description": "Search the web", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - } - ], - stream=True, - ) - - # Response should be a streaming iterator - chunks = [] - async for chunk in response: - chunks.append(chunk) - - # Verify we got streaming chunks - assert len(chunks) > 0 - - # Verify chunks have expected structure - for chunk in chunks: - assert hasattr(chunk, "choices") - assert len(chunk.choices) > 0 - - finally: - litellm.callbacks = [] - - @pytest.mark.asyncio async def test_maybe_run_chat_completion_agentic_loop_calls_chat_completion_hook(): """Regression test: maybe_run_chat_completion_agentic_loop must call diff --git a/tests/unit/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py index 2bb93a58531..5d4c9e65d9b 100644 --- a/tests/unit/litellm_core_utils/test_logging_worker.py +++ b/tests/unit/litellm_core_utils/test_logging_worker.py @@ -180,6 +180,39 @@ class TestLoggingWorker: assert sorted(fired) == ["first", "second"] + def test_callback_finishing_after_loop_change_settles_only_its_own_queue(self): + worker = LoggingWorker(timeout=1.0, max_queue_size=10, concurrency=1) + fired = [] + + async def marker(name, delay=0.0): + await asyncio.sleep(delay) + fired.append(name) + + async def start_slow_callback(): + worker.ensure_initialized_and_enqueue(marker("slow", delay=0.05)) + await asyncio.sleep(0.01) + + async def log_on_second_loop(): + for name in ("b1", "b2", "b3"): + worker.ensure_initialized_and_enqueue(marker(name)) + for _ in range(2): + await asyncio.sleep(0) + + first_loop = asyncio.new_event_loop() + try: + first_loop.run_until_complete(start_slow_callback()) + first_loop_tasks = tuple(asyncio.all_tasks(first_loop)) + asyncio.run(log_on_second_loop()) + first_loop.run_until_complete(asyncio.sleep(0.1)) + failures = [ + task.exception() for task in first_loop_tasks if task.done() and not task.cancelled() and task.exception() + ] + finally: + first_loop.close() + + assert failures == [] + assert "slow" in fired + @pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"]) def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded): """ diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 75d3a23e012..f7ded4f3fa8 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -29,6 +29,7 @@ import litellm.constants from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( + _encoding_count, _get_exact_count_function, _get_extrapolating_count_function, _get_tiktoken_count_function, @@ -79,15 +80,17 @@ def test_token_counter_basic(): ) -def test_token_counter_large_repeated_text_is_fast(): - messages = [{"role": "user", "content": [{"type": "text", "text": "A" * 1024 * 1024}]}] +def test_token_counter_large_repeated_text_is_encoded_in_bounded_chunks(): + text_length: Final = 1024 * 1024 + messages: Final = [{"role": "user", "content": [{"type": "text", "text": "A" * text_length}]}] - start_time = time.perf_counter() - tokens = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages) - elapsed = time.perf_counter() - start_time + with patch("litellm.litellm_core_utils.token_counter._encoding_count", wraps=_encoding_count) as encoding_count: + tokens: Final = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages) - assert elapsed < 2, f"Token counting took too long: {elapsed:.2f}s" + encoded_lengths: Final = tuple(len(call.args[1]) for call in encoding_count.call_args_list) assert tokens > 0 + assert sum(encoded_lengths) >= text_length + assert max(encoded_lengths) <= litellm.constants.TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS @pytest.mark.parametrize( diff --git a/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py index 179a6cad4aa..138ad8bb81f 100644 --- a/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py +++ b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py @@ -222,7 +222,8 @@ def test_langflow_extra_body_cannot_inject_tweaks_into_run_payload(): def fake_post(*args, **kwargs): body = kwargs.get("data") - posted_bodies.append(json.loads(body) if isinstance(body, str) else body) + if str(kwargs.get("url", "")).startswith("http://example.com"): + posted_bodies.append(json.loads(body) if isinstance(body, (str, bytes)) else body) resp = MagicMock(spec=httpx.Response) resp.status_code = 200 resp.json.return_value = {"outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}]} diff --git a/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py index 6115e627b26..fb5e84c8294 100644 --- a/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py +++ b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py @@ -3717,7 +3717,7 @@ async def test_auth_vertex_ai_route(prisma_client): @pytest.mark.asyncio -async def test_user_api_key_auth_db_unavailable(): +async def test_user_api_key_auth_db_unavailable(monkeypatch): """ Test that user_api_key_auth handles DB connection failures appropriately when: 1. DB connection fails during token validation @@ -3747,7 +3747,7 @@ async def test_user_api_key_auth_db_unavailable(): # Set up test environment setattr(litellm.proxy.proxy_server, "prisma_client", MockPrismaClient()) - setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache()) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache()) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr( litellm.proxy.proxy_server, @@ -3777,7 +3777,7 @@ async def test_user_api_key_auth_db_unavailable(): @pytest.mark.asyncio -async def test_user_api_key_auth_db_unavailable_not_allowed(): +async def test_user_api_key_auth_db_unavailable_not_allowed(monkeypatch): """ Test that user_api_key_auth raises an exception when: This is default behavior @@ -3808,7 +3808,7 @@ async def test_user_api_key_auth_db_unavailable_not_allowed(): # Set up test environment setattr(litellm.proxy.proxy_server, "prisma_client", MockPrismaClient()) - setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache()) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache()) setattr(litellm.proxy.proxy_server, "general_settings", {}) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 65b368ca9e3..8947da4d9fc 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -1,5 +1,6 @@ import os import traceback +from typing import Final from unittest import mock from dotenv import load_dotenv @@ -35,6 +36,7 @@ from fastapi import FastAPI from fastapi.testclient import TestClient from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.proxy_server import ( # Replace with the actual module where your FastAPI router is defined app, initialize, @@ -418,7 +420,7 @@ def test_chat_completion_forward_llm_provider_auth_headers( @mock_patch_acompletion() @pytest.mark.asyncio -async def test_team_disable_guardrails(mock_acompletion, client_no_auth): +async def test_team_disable_guardrails(mock_acompletion, client_no_auth, monkeypatch): """ If team not allowed to turn on/off guardrails @@ -438,8 +440,9 @@ async def test_team_disable_guardrails(mock_acompletion, client_no_auth): UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.proxy_server import hash_token, user_api_key_cache + from litellm.proxy.proxy_server import hash_token + user_api_key_cache: Final = UserApiKeyCache() _team_id = "1234" user_key = "sk-12345678" @@ -459,7 +462,7 @@ async def test_team_disable_guardrails(mock_acompletion, client_no_auth): user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token) user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj) - setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world") @@ -481,10 +484,11 @@ from tests.unit.proxy.test_custom_callback_input import CompletionCustomHandler @mock_patch_acompletion() -def test_custom_logger_failure_handler(mock_acompletion, client_no_auth): +def test_custom_logger_failure_handler(mock_acompletion, client_no_auth, monkeypatch): from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.proxy_server import hash_token, user_api_key_cache + from litellm.proxy.proxy_server import hash_token + user_api_key_cache: Final = UserApiKeyCache() rpm_limit = 0 mock_api_key = "sk-my-test-key" @@ -501,7 +505,7 @@ def test_custom_logger_failure_handler(mock_acompletion, client_no_auth): litellm.callbacks = [mock_logger, mock_logger_unit_tests] proxy_logging_obj._init_litellm_callbacks(llm_router=None) - setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR") setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj) @@ -1296,7 +1300,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa @pytest.mark.parametrize("team_route", ["/team/member_add", "/team/member_delete"]) @pytest.mark.asyncio async def test_create_team_member_add_team_admin_user_api_key_auth( - prisma_client, team_member_role, team_route # noqa: F811 # pytest fixture, not a redefinition + prisma_client, team_member_role, team_route, monkeypatch # noqa: F811 # pytest fixture, not a redefinition ): import time @@ -1307,9 +1311,10 @@ async def test_create_team_member_add_team_admin_user_api_key_auth( ProxyException, hash_token, user_api_key_auth, - user_api_key_cache, ) + user_api_key_cache: Final = UserApiKeyCache() + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm, "max_internal_user_budget", 10) @@ -1335,7 +1340,7 @@ async def test_create_team_member_add_team_admin_user_api_key_auth( user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj) - setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) ## TEST IF TEAM ADMIN ALLOWED TO CALL /MEMBER_ADD ENDPOINT import json @@ -2349,7 +2354,7 @@ async def test_proxy_server_prisma_setup(): mock_client.db = mock_db prisma_client = await ProxyStartupEvent._setup_prisma_client( - database_url=os.getenv("DATABASE_URL"), + database_url="postgresql://user:pass@localhost:5432/litellm", proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), user_api_key_cache=user_api_key_cache, ) diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 3401a335b2f..2d66524326f 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -14120,6 +14120,7 @@ class TestContextWindowEscalation: assert result.routing_decision["context_escalated"] is True @pytest.mark.asyncio + @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize( "deployments,tiers,expected_model", [ diff --git a/tests/unit/secret_managers/test_cyberark_secret_manager.py b/tests/unit/secret_managers/test_cyberark_secret_manager.py index 3f3669ab9ef..334e6437f1d 100644 --- a/tests/unit/secret_managers/test_cyberark_secret_manager.py +++ b/tests/unit/secret_managers/test_cyberark_secret_manager.py @@ -1,7 +1,9 @@ +import asyncio import json from pathlib import Path from typing import Final, TypedDict, cast +import httpx import pytest import respx @@ -107,6 +109,86 @@ async def test_async_write_matches_parity_fixture(monkeypatch: pytest.MonkeyPatc assert value_route.calls.last.request.content == b"v" +@pytest.mark.asyncio +@respx.mock +async def test_async_write_retries_policy_load_conflict(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + fixture: Final = _fixture() + manager: Final = _configure_manager(monkeypatch, fixture) + secret: Final = fixture["secrets"][0] + endpoint: Final = fixture["endpoint"] + _respond(respx.post(endpoint + fixture["authenticate_path"]), content=fixture["token_json"].encode()) + policy_route: Final = respx.post(endpoint + fixture["policy_path"]).mock( + side_effect=[httpx.Response(409), httpx.Response(409), httpx.Response(201)] + ) + value_route: Final = respx.post(endpoint + secret["path"]).mock( + side_effect=lambda _: httpx.Response(201 if policy_route.call_count == 3 else 404) + ) + + result: Final = await manager.async_write_secret(secret["name"], "v") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy secret manager API is untyped + + assert policy_route.call_count == 3 + assert value_route.call_count == 1 + assert result["status"] == "success" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "policy_outcome", + [422, 500, httpx.ConnectError("conjur unreachable")], + ids=["unprocessable", "server_error", "unreachable"], +) +@respx.mock +async def test_async_write_does_not_retry_non_conflict_policy_failures( + monkeypatch: pytest.MonkeyPatch, policy_outcome: int | httpx.ConnectError +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + fixture: Final = _fixture() + manager: Final = _configure_manager(monkeypatch, fixture) + secret: Final = fixture["secrets"][0] + endpoint: Final = fixture["endpoint"] + _respond(respx.post(endpoint + fixture["authenticate_path"]), content=fixture["token_json"].encode()) + policy_route: Final = respx.post(endpoint + fixture["policy_path"]) + if isinstance(policy_outcome, int): + _respond(policy_route, status_code=policy_outcome) + else: + policy_route.mock(side_effect=policy_outcome) + value_route: Final = _respond(respx.post(endpoint + secret["path"]), status_code=201) + + await manager.async_write_secret(secret["name"], "v") # pyright: ignore[reportUnknownMemberType] # legacy secret manager API is untyped + + assert policy_route.call_count == 1 + assert value_route.call_count == 1 + + +@pytest.mark.asyncio +@respx.mock +async def test_concurrent_async_writes_load_policy_one_at_a_time(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + fixture: Final = _fixture() + manager: Final = _configure_manager(monkeypatch, fixture) + endpoint: Final = fixture["endpoint"] + _respond(respx.post(endpoint + fixture["authenticate_path"]), content=fixture["token_json"].encode()) + in_flight: Final = asyncio.Semaphore(1) + + async def load_policy(_: httpx.Request) -> httpx.Response: + if in_flight.locked(): + return httpx.Response(409) + async with in_flight: + await asyncio.sleep(0.05) + return httpx.Response(201) + + policy_route: Final = respx.post(endpoint + fixture["policy_path"]).mock(side_effect=load_policy) + respx.post(url__startswith=endpoint + "/secrets/").respond(status_code=201) # pyright: ignore[reportUnknownMemberType] # respx route stubs leave response builder partially unknown + + results: Final = await asyncio.gather( + *(manager.async_write_secret(f"concurrent-{index}", "v") for index in range(4)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy secret manager API is untyped + ) + + assert policy_route.call_count == 4 + assert [result["status"] for result in results] == ["success"] * 4 + + def test_missing_credentials_raise_value_error(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) for name in ( diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index d9cfe88d52f..5cdecc31575 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -9,7 +9,7 @@ import sys import time from io import StringIO from pathlib import Path -from typing import List +from typing import Final, List import pytest from pydantic import BaseModel, computed_field @@ -1584,6 +1584,8 @@ def _emit_access_line(full_path: str) -> str: handler = logging.StreamHandler(stream) handler.setFormatter(AccessFormatter('%(client_addr)s - "%(request_line)s" %(status_code)s', use_colors=False)) saved_level, saved_propagate = logger.level, logger.propagate + saved_filters: Final = logger.filters[:] + logger.filters = [f for f in saved_filters if type(f).__module__.split(".")[0] == "litellm"] logger.addHandler(handler) logger.setLevel(logging.INFO) logger.propagate = False @@ -1593,6 +1595,7 @@ def _emit_access_line(full_path: str) -> str: logger.removeHandler(handler) logger.setLevel(saved_level) logger.propagate = saved_propagate + logger.filters = saved_filters return stream.getvalue() diff --git a/tests/unit/test_video_generation.py b/tests/unit/test_video_generation.py index 644c7a41f49..5c1d0bfa884 100644 --- a/tests/unit/test_video_generation.py +++ b/tests/unit/test_video_generation.py @@ -11,6 +11,7 @@ import litellm from litellm.cost_calculator import default_video_cost_calculator from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.gemini.videos.transformation import GeminiVideoConfig @@ -988,6 +989,7 @@ class TestVideoLogging: """ custom_logger = self.TestVideoLogger() litellm.logging_callback_manager._reset_all_callbacks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) litellm.callbacks = [custom_logger] # Mock video generation response