litellm/scripts/test_tool_allowlist_script.py

#!/usr/bin/env python3
"""
Standalone script to test tool allowlist enforcement and tool name extraction.

Run from repo root:
  poetry run python scripts/test_tool_allowlist_script.py

Or run the unit tests:
  poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v
"""

import asyncio
import sys
from pathlib import Path

# Ensure repo root is on path
repo_root = Path(__file__).resolve().parent.parent
if str(repo_root) not in sys.path:
    sys.path.insert(0, str(repo_root))


def test_extraction():
    """Test extract_request_tool_names for each API shape."""
    from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names

    cases = [
        ("OpenAI chat tools", "/v1/chat/completions", {"tools": [{"type": "function", "function": {"name": "get_weather"}}]}),
        ("OpenAI chat functions", "/v1/chat/completions", {"functions": [{"name": "run_sql"}]}),
        ("OpenAI responses function", "/v1/responses", {"tools": [{"type": "function", "name": "get_current_weather"}]}),
        ("OpenAI responses MCP", "/v1/responses", {"tools": [{"type": "mcp", "server_label": "dmcp"}]}),
        ("Anthropic", "/v1/messages", {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}),
        ("Google generateContent", "/generate_content", {"tools": [{"functionDeclarations": [{"name": "schedule_meeting"}]}]}),
        ("MCP call_tool", "/mcp/call_tool", {"name": "my_tool", "arguments": {}}),
        ("Non-tool route", "/v1/embeddings", {"tools": [{"type": "function", "function": {"name": "x"}}]}),
    ]
    print("=== extract_request_tool_names(route, data) ===\n")
    for label, route, data in cases:
        names = extract_request_tool_names(route, data)
        print(f"  {label}: {names}")
    print()


async def test_check_tools_allowlist():
    """Test check_tools_allowlist with mock tokens."""
    from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
    from litellm.proxy.auth.auth_checks import check_tools_allowlist

    def token(metadata=None, team_metadata=None):
        return UserAPIKeyAuth(
            api_key="test-key",
            user_id="user",
            team_id="team",
            org_id=None,
            models=["*"],
            metadata=metadata or {},
            team_metadata=team_metadata or {},
        )

    print("=== check_tools_allowlist (auth) ===\n")

    # No allowlist -> pass
    await check_tools_allowlist(
        request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
        valid_token=token(),
        team_object=None,
        route="/v1/chat/completions",
    )
    print("  No allowlist, body has tools: PASS")

    # Allowed tool -> pass
    await check_tools_allowlist(
        request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
        valid_token=token(metadata={"allowed_tools": ["get_weather"]}),
        team_object=None,
        route="/v1/chat/completions",
    )
    print("  allowed_tools=['get_weather'], body has get_weather: PASS")

    # Disallowed tool -> raise
    try:
        await check_tools_allowlist(
            request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
            valid_token=token(metadata={"allowed_tools": ["other_tool"]}),
            team_object=None,
            route="/v1/chat/completions",
        )
        print("  DISALLOWED: expected ProxyException")
    except ProxyException as e:
        if e.type == ProxyErrorTypes.tool_access_denied:
            print("  allowed_tools=['other_tool'], body has get_weather: PASS (raised tool_access_denied)")
        else:
            print(f"  Unexpected ProxyException type: {e.type}")
    except Exception as e:
        print(f"  Unexpected: {e}")

    # Team allowlist when key empty
    await check_tools_allowlist(
        request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
        valid_token=token(team_metadata={"allowed_tools": ["get_weather"]}),
        team_object=None,
        route="/v1/chat/completions",
    )
    print("  team_metadata.allowed_tools=['get_weather']: PASS")
    print()


def main():
    print("Tool allowlist / tool name extraction – script checks\n")
    test_extraction()
    asyncio.run(test_check_tools_allowlist())
    print("Done. For full unit tests run:")
    print("  poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v")


if __name__ == "__main__":
    main()