Files
litellm/tests/test_litellm/test_ssl_verify_unit.py
T
Sameer Kankute e8fcb01215 Litellm OSS Staging (#29161)
* Cato Networks guardrail, based on Aim (#26597)

* Aim was acquired by Cato Networks, creating Cato Networks guardrail based on Aim

* Add more tests

* Move test so they are reached by codecov coverage

* base URL trailing slashes

* Support Lemonade runtime context metadata (#28135)

* Support Lemonade runtime context metadata

* Add provider hook for runtime model metadata

* Address provider model info review feedback

Keep the runtime model info hook duck-typed instead of extending the base model-info class, and avoid importing ModelInfoBase from Ollama common utilities to reduce CodeQL cyclic-import noise.

Co-authored-by: openhands <openhands@all-hands.dev>

* Fix CI after staging rebase

Relax the Ollama runtime metadata return annotation to match the provider-hook dict response and update the Google Interactions OpenAPI status expectation for the current live spec.

Co-authored-by: openhands <openhands@all-hands.dev>

* Normalize Lemonade runtime model metadata

* Avoid leaking Ollama metadata auth

* Avoid leaking Lemonade metadata auth

---------

Co-authored-by: Graham Neubig <398875+neubig@users.noreply.github.com>
Co-authored-by: openhands <openhands@all-hands.dev>

* fix(cato): address guardrail review feedback

Use proxy-authenticated user identity, forward moderation hook return values,
and ensure streaming sender tasks are cancelled and awaited on exit.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path - clone of #28010 (#28846)

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path

Fixes #26083

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the
NON_GEMINI route. Per owtaylor's plan on #26083: add the google/gemma-
prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up
and should_use_openai_handler routes it to the OpenAI-compatible
/endpoints/openapi/chat/completions URL. No gemma-detection exclusion
needed (the "gemma/" check uses a slash, which google/gemma-... doesn't
match). No OpenAIGPTConfig subclass needed — works with the base handler.

* fix(vertex_ai): mark gemma-4-26b-a4b-it-maas as vision-capable (empirically verified)

* fix(vertex_ai): address greptile feedback — provider category, canonical URL, sync backup

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

   Addresses oss-pr-review-agent-shin feedback on PR #28010:
   supports_function_calling, supports_tool_choice, and supports_vision were
   marked true but had no tests proving the payloads actually reached the
   OpenAI-compatible endpoint.

   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: Delete uv.lock

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

Addresses oss-pr-review-agent-shin feedback on PR #28010:

   P1 (patch target): Added a comment explaining why patching
   litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct —
   get_async_httpx_client() (defined in http_handler.py) instantiates
   AsyncHTTPHandler within that module's scope, so the definition-site patch
   intercepts it. Without the mock the test raises AuthenticationError,
   confirming it never silently passes.

   P2 (partner-provider regression guard): Added
   test_gemma_routes_through_openai_handler() which calls
   VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's
   routing to VertexPartnerProvider.llama ever changes the URL-shape tests
   below it become a real regression guard rather than an unanchored unit test.

   Also added:
   - test_gemma_maas_supports_function_calling / supports_vision — capability
     flag checks via patch.dict(litellm.model_cost)
   - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice
     forwarded in the request body
   - test_vertex_ai_gemma_vision_passthrough — image_url part survives
     transformation to the global endpoint
   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: proper patch for unit tests

---------

Co-authored-by: Iana <iana@Shivakumars-MacBook-Pro.local>

* fix(cato): guardrail all completion choices on output

When n > 1, only choices[0] was analyzed and redacted. Iterate every
Choices entry so block and anonymize actions apply to all completions.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix review

* fix(cato_networks): harden output anonymize handling and restructure nested UI routes

Guard against empty redacted_output and empty all_redacted_messages from Cato.
Restructure nested admin UI HTML exports to index.html so extensionless routes work.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix mypy

* fix(cato): guard missing policy_drill_down and all_redacted_messages keys

* fix(cato): avoid KeyError bypassing block action on missing analysis_result

* fix(cato): preserve non-text message fields during anonymize

Rebuild redacted messages from the original messages, overwriting only
content, so tool_calls, tool_call_id, name and multimodal fields survive
the anonymize action.

* fix(cato): preserve trailing messages when fewer redacted messages returned

Avoid silently truncating the conversation in _anonymize_request when Cato
returns fewer redacted messages than were sent, and isolate the no-api-key
config test from a pre-existing CATO_API_KEY environment variable.

* fix(cato,model-info): preserve stream block signal on sender teardown; forward api_key in dynamic model-info lookup

Suppress ConnectionClosed (alongside CancelledError) when tearing down the
Cato streaming sender task so a backend ConnectionClosed cannot mask the
original StreamingCallbackError (e.g. a guardrail block) raised by the
receive loop.

Thread api_key through get_model_info -> _get_model_info_helper so an
explicit key reaches a provider's dynamic get_model_info for a caller-supplied
api_base. Previously only api_base was forwarded, so authenticated Ollama and
Lemonade servers at a custom base could only be queried unauthenticated.

* fix(cato): surface mid-stream forwarding errors instead of blocking on recv

If the upstream LLM stream errors mid-flight, the sender task dies before
sending the terminal done frame, so the consumer would block on websocket.recv()
until Cato closes the connection. Race recv against the sender task and raise the
stored sender exception promptly as a StreamingCallbackError.

* fix(cato): drop spoofable end_user_id from guardrail user identity

Only the key/JWT-bound user_email is a trusted identity. end_user_id is
resolved from caller-supplied request fields (OpenAI user param, headers,
metadata), so an authenticated caller with no bound user_email could set it
to another user's email and have LiteLLM forward x-cato-user-email for that
victim, poisoning Cato audit and policy attribution. Forward only user_email
and omit the header otherwise.

* fix(cato): harden output anonymize path against missing content key

* fix(cato): fall back to original message when redacted content key is missing

* refactor(model-info): drop unused api_key from cached model-info helper

_cached_get_model_info_helper is only called by the cost-tracking hot path,
which never authenticates, so the api_key parameter was never populated.
Keeping it in the lru_cache key offered no benefit and risked fragmenting
the high-RPS cache and retaining credential strings per entry.

* fix(cato): preserve None content on tool-call-only choices in output hook

* fix(ollama): respect static-model guard in OllamaConfig.get_model_info

Delegate to OllamaModelInfo.get_model_info so statically-priced Ollama
models short-circuit before the /api/show network call instead of
hitting the server unconditionally.

* fix(lemonade,ollama): treat empty api_key as unset to avoid leaking server creds

An empty-string api_key was treated as an explicit key, so it passed the
guard meant to keep server-side credentials off caller-supplied bases and
then fell back through the env/global key chain. A caller could point
api_base at a server they control and send api_key="" to receive the
configured provider key in the Authorization header. Gate the credential
fallback on the api_key being truthy instead of merely not-None.

* fix(cato): inspect and redact Responses-API input, not just messages

The guardrail only read data["messages"], so /v1/responses requests, which
carry their text in data["input"], reached Cato as an empty message list
and bypassed inspection entirely. Send build_inspection_messages(data) so
both shapes are analyzed, and write anonymized results back with
apply_redacted_messages_back when the request used input.

* perf(utils): keep api_key out of get_model_info lru_cache key

* fix(cato): propagate ssl_verify to streaming WebSocket connection

The streaming hook applied ssl_verify only to the HTTP handler; the
websockets.connect() call used default verification, so a custom Cato
instance behind TLS with a self-signed cert worked for non-streaming
calls but failed every streaming request. Resolve the ssl_verify setting
into the connect() ssl argument, mirroring the HTTP handler.

* refactor(utils): rename shadowing local in _get_model_info_helper

* fix(cato): flatten multimodal chat content before inspection

Chat Completions requests whose message content is a multimodal parts
array were posted to Cato as the raw OpenAI parts, so text inside
content: [{"type":"text", ...}] reached the model without Cato ever
inspecting the string. Flatten each message's list content to plain text
while keeping the list 1:1 with the request so the index-based redaction
write-back stays valid; Responses-API input requests still go through
build_inspection_messages.

* test(lemonade): clear get_model_info cache around api_base test

* fix(cato): inspect and redact Responses-API input even when messages present

_inspection_messages returned early once messages was non-empty, so a
/v1/responses caller could place benign text in messages and disallowed
text in input and have only messages reach Cato while the model used
input. Inspect both fields and write anonymize redactions back to input
as well as the index-aligned messages.

* test(log_db_metrics): assert table_name event_metadata contract

log_db_metrics now emits minimal event_metadata via _safe_db_event_metadata
(table_name only, function_name/function_kwargs/function_args dropped as
redundant with call_type and unsafe to stamp on a span). The success-path
test still asserted function_name membership and crashed with TypeError on
the None metadata returned when no table_name is passed. Pass a table_name
and assert the surfaced contract instead.

* fix(cato): inspect and redact completion prompt and Responses-API instructions

The Cato guardrail only inspected chat messages and the Responses-API input field, so blocked text placed in the legacy /v1/completions prompt or the /v1/responses instructions field reached the model without ever being sent to Cato. Both fields are now appended as synthetic inspection messages, and the anonymize path slices Cato's redactions back to the field they came from.

* fix(cato): serialize non-str/bytes websocket chunks before forwarding

* fix(cato): inspect tool descriptions and tool-call arguments

* fix(cato): map redacted output by assistant index; restore get_model_info.cache_info

* fix(cato): block output even when detection_message is null/empty

A block_action returned by Cato on the output hook whose detection_message
was null or empty was let through to the caller: the truthiness guard on
detection_message skipped the HTTPException and the unblocked response was
returned. Raise the HTTPException directly in _handle_block_action_on_output
so the output path blocks unconditionally, mirroring the input path.

* fix(cato): inspect and redact nested tool param and legacy function descriptions

Tool/function parameter descriptions and the legacy functions[] array are
forwarded to the model but were not seen by Cato, so blocked text hidden there
bypassed inspection and anonymization. Recursively walk every description string
in tools[].function and functions[] schemas for both the analyze payload and the
anonymize write-back.

* fix(cato): traverse schema descriptions iteratively to satisfy recursive detector

The nested walk() generator recursed over tool/function JSON schemas with no
depth bound, which the recursive_detector code-quality gate rejects. Replace it
with an explicit-stack DFS that yields the same (container, key) refs in the
same pre-order, so schema description redaction is unchanged.

* fix(cato): inspect and redact response_format JSON schema descriptions

response_format json_schema descriptions are forwarded to the model, so
blocked text hidden in nested schema descriptions could bypass Cato
inspection and redaction. Extend the schema-description walk to cover
response_format alongside tools and legacy functions.

* fix(cato): skip output rewrite when Cato returns no redaction

Return None from call_cato_guardrail_on_output on monitor/no-action so the
post-call hook only mutates the message when there is an actual redaction,
instead of redundantly re-writing the original content.

* refactor(utils): resolve explicit api_key model info without the cache

Move the model-info build into a non-cached _build_model_info helper and drop
api_key from the lru-cached _cached_get_model_info signature. Both cached
helpers now take the same (model, provider, api_base) key and never forward
api_key, while explicit per-caller keys are resolved through the builder
directly instead of reaching into the cache wrapper's __wrapped__.

* fix(cato): inspect and redact non-description schema string values

Tool, function and response_format JSON schemas forward more than just
description text to the model. enum, const, default, examples and title
values are sent verbatim, so blocked content hidden in any of them
bypassed Cato inspection and redaction. Walk those schema string values
alongside descriptions on both the inspection and anonymize paths.

* fix(model-info): surface swallowed dynamic model-info errors

The provider-specific get_model_info dispatch falls back to the static cost
map when a provider's dynamic lookup raises, which is intentional graceful
degradation. Previously the exception was discarded with a bare debug line,
so a real failure (e.g. a provider whose get_model_info signature does not
accept api_key) was invisible. Log the exception at warning level with the
model and provider context so the fallback is diagnosable.

* fix(cato): inspect and redact Responses API output in post-call hook

The post-call success hook only handled ModelResponse, so /v1/responses
(which returns a ResponsesAPIResponse) bypassed the Cato output guardrail.
Extract and inspect/redact every output_text content block and function-call
arguments string, blocking on a block action, so generated text cannot escape
inspection by using the Responses API.

* chore: reset _experimental/out folder

* chore(ui): remove orphaned prebuilt dashboard chunk files

The _experimental/out manifests are byte-identical to the base branch, so the
served dashboard already matches base. 436 unreferenced Next.js chunk files had
accumulated in the directory and are not loaded by any manifest; removing them
restores the committed UI artifacts to the base build and drops the artifact
churn from this PR's diff.

* fix(guardrails,ollama): forward ssl_verify to Cato init and raise_for_status on /api/show

---------

Co-authored-by: Alex Yaroslavsky <trexinc@gmail.com>
Co-authored-by: Graham Neubig <neubig@gmail.com>
Co-authored-by: Graham Neubig <398875+neubig@users.noreply.github.com>
Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Piotr Placzko <piotr@icep-design.com>
Co-authored-by: Iana <iana@Shivakumars-MacBook-Pro.local>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-06-01 21:22:35 -07:00

236 lines
8.9 KiB
Python

"""
Unit tests for per-service SSL support in LiteLLM.
These tests verify that ssl_verify parameters are correctly propagated
through the call stack without requiring live API credentials.
"""
import sys
from pathlib import Path
from unittest.mock import Mock, patch
import pytest
# Add litellm to path
sys.path.insert(0, str(Path(__file__).parent))
import litellm.proxy.guardrails.guardrail_hooks.aim.aim as _aim_module
import litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks as _cato_networks_module
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
from litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks import CatoNetworksGuardrail
class TestBaseAWSLLMSSLVerify:
"""Test SSL verification parameter handling in BaseAWSLLM."""
def test_get_ssl_verify_with_parameter(self):
"""Test that _get_ssl_verify accepts and uses the ssl_verify parameter."""
base_llm = BaseAWSLLM()
# Test with True
result = base_llm._get_ssl_verify(ssl_verify=True)
assert result is True
# Test with False
result = base_llm._get_ssl_verify(ssl_verify=False)
assert result is False
# Test with cert path
cert_path = "/path/to/cert.pem"
result = base_llm._get_ssl_verify(ssl_verify=cert_path)
assert result == cert_path
def test_get_ssl_verify_without_parameter(self):
"""Test that _get_ssl_verify falls back to environment/global when no parameter."""
base_llm = BaseAWSLLM()
# Should fall back to environment or global litellm.ssl_verify
result = base_llm._get_ssl_verify()
# Result depends on environment, just verify it doesn't crash
assert result is not None or result is None # Can be None, True, False, or path
@patch("boto3.client")
def test_get_credentials_propagates_ssl_verify(self, mock_boto_client):
"""Test that get_credentials propagates ssl_verify to boto3 clients."""
base_llm = BaseAWSLLM()
# Mock the boto3 client
mock_sts_client = Mock()
mock_sts_client.assume_role.return_value = {
"Credentials": {
"AccessKeyId": "test_key",
"SecretAccessKey": "test_secret",
"SessionToken": "test_token",
"Expiration": "2026-01-20T00:00:00Z",
}
}
mock_boto_client.return_value = mock_sts_client
# Call get_credentials with ssl_verify parameter
cert_path = "/path/to/cert.pem"
try:
base_llm.get_credentials(
aws_access_key_id="test_key",
aws_secret_access_key="test_secret",
aws_region_name="us-east-1",
ssl_verify=cert_path,
)
except Exception:
# May fail due to missing credentials, but we're checking the call
pass
# Verify boto3.client was called with verify parameter
# Note: This test verifies the parameter is accepted, actual propagation
# is tested in integration tests
assert True # If we got here without error, parameter was accepted
class TestBedrockLLMSSLVerify:
"""Test SSL verification parameter handling in BedrockLLM."""
def test_bedrock_llm_accepts_ssl_verify_in_optional_params(self):
"""Test that BedrockLLM can receive ssl_verify in optional_params."""
# This is a simple test to verify the parameter is accepted
# The actual propagation is tested in integration tests
bedrock_llm = BedrockLLM()
# Verify the class exists and can be instantiated
assert bedrock_llm is not None
# Verify _get_ssl_verify method exists and works
result = bedrock_llm._get_ssl_verify(ssl_verify="/path/to/cert.pem")
assert result == "/path/to/cert.pem"
class TestAimGuardrailSSLVerify:
"""Test SSL verification parameter handling in AimGuardrail."""
def test_init_accepts_ssl_verify(self):
"""Test that AimGuardrail.__init__ accepts and uses ssl_verify parameter."""
mock_handler = Mock()
# Use patch.object on the actual module reference for reliable patching
# across different import orders / CI environments
with patch.object(
_aim_module, "get_async_httpx_client", return_value=mock_handler
) as mock_get_client:
# Initialize with ssl_verify
cert_path = "/path/to/aim_cert.pem"
AimGuardrail(
api_key="test_key",
api_base="https://test.aim.api",
ssl_verify=cert_path,
)
# Verify get_async_httpx_client was called with ssl_verify in params
assert mock_get_client.called
call_kwargs = mock_get_client.call_args[1]
assert "params" in call_kwargs
assert call_kwargs["params"] is not None
assert call_kwargs["params"]["ssl_verify"] == cert_path
def test_init_without_ssl_verify(self):
"""Test that AimGuardrail works without ssl_verify parameter."""
mock_handler = Mock()
# Use patch.object on the actual module reference for reliable patching
with patch.object(
_aim_module, "get_async_httpx_client", return_value=mock_handler
) as mock_get_client:
# Initialize without ssl_verify
AimGuardrail(api_key="test_key", api_base="https://test.aim.api")
# Should still work, just without custom SSL
assert mock_get_client.called
class TestCatoNetworksGuardrailSSLVerify:
"""Test SSL verification parameter handling in CatoNetworksGuardrail."""
def test_init_accepts_ssl_verify(self):
"""Test that CatoNetworksGuardrail.__init__ accepts and uses ssl_verify parameter."""
mock_handler = Mock()
# Use patch.object on the actual module reference for reliable patching
# across different import orders / CI environments
with patch.object(
_cato_networks_module, "get_async_httpx_client", return_value=mock_handler
) as mock_get_client:
# Initialize with ssl_verify
cert_path = "/path/to/cato_cert.pem"
CatoNetworksGuardrail(
api_key="test_key",
api_base="https://test.catonetworks.api",
ssl_verify=cert_path,
)
# Verify get_async_httpx_client was called with ssl_verify in params
assert mock_get_client.called
call_kwargs = mock_get_client.call_args[1]
assert "params" in call_kwargs
assert call_kwargs["params"] is not None
assert call_kwargs["params"]["ssl_verify"] == cert_path
def test_init_without_ssl_verify(self):
"""Test that CatoNetworksGuardrail works without ssl_verify parameter."""
mock_handler = Mock()
# Use patch.object on the actual module reference for reliable patching
with patch.object(
_cato_networks_module, "get_async_httpx_client", return_value=mock_handler
) as mock_get_client:
# Initialize without ssl_verify
CatoNetworksGuardrail(api_key="test_key", api_base="https://test.catonetworks.api")
# Should still work, just without custom SSL
assert mock_get_client.called
class TestHTTPHandlerSSLVerify:
"""Test SSL verification parameter handling in HTTP handlers."""
def test_get_async_httpx_client_accepts_ssl_verify_in_params(self):
"""Test that get_async_httpx_client accepts ssl_verify in params dict."""
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
# Call with ssl_verify in params
cert_path = "/path/to/cert.pem"
client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
params={"ssl_verify": cert_path},
)
# Verify client was created (actual SSL config is tested in integration tests)
assert client is not None
def test_ssl_verify_parameter_types():
"""Test that various ssl_verify parameter types are handled correctly."""
base_llm = BaseAWSLLM()
# Test boolean True
result = base_llm._get_ssl_verify(ssl_verify=True)
assert result is True
# Test boolean False
result = base_llm._get_ssl_verify(ssl_verify=False)
assert result is False
# Test string path
cert_path = "/path/to/cert.pem"
result = base_llm._get_ssl_verify(ssl_verify=cert_path)
assert result == cert_path
# Test None (should fall back to environment/global)
result = base_llm._get_ssl_verify(ssl_verify=None)
# Result depends on environment
assert result is not None or result is None
if __name__ == "__main__":
# Run tests
pytest.main([__file__, "-v", "--tb=short"])