diff --git a/litellm/router.py b/litellm/router.py index 585cec682d4..6c179fa6948 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -14,6 +14,7 @@ import hashlib import inspect import json import logging +import re import threading import time import traceback @@ -1551,7 +1552,9 @@ class Router: # Drain any fire-and-forget tasks (e.g. alerting hooks) # scheduled via asyncio.create_task during acompletion. pending = asyncio.all_tasks() - pending.discard(asyncio.current_task()) + current = asyncio.current_task() + if current is not None: + pending.discard(current) if pending: await asyncio.gather(*pending, return_exceptions=True) @@ -6526,6 +6529,18 @@ class Router: deployment = self._add_deployment(deployment=deployment) + # Validate tag_regex patterns early so a bad regex fails at startup + # rather than silently misbehaving on the first matching request. + _tag_regex = deployment.litellm_params.get("tag_regex") or [] + for pattern in _tag_regex: + try: + re.compile(pattern) + except re.error as exc: + raise ValueError( + f"Invalid regex in tag_regex for model '{deployment.model_name}': " + f"{pattern!r} — {exc}" + ) from exc + model = deployment.to_json(exclude_none=True) self._add_model_to_list_and_index_map( diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index e7156bf1282..d7e756f3d69 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -6,6 +6,7 @@ Use this to route requests between Teams - If no default_deployments are set, return all deployments """ +import re from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union from litellm._logging import verbose_logger @@ -19,6 +20,29 @@ else: LitellmRouter = Any +def _is_valid_deployment_tag_regex( + tag_regexes: List[str], + header_strings: List[str], +) -> Optional[str]: + """ + Test compiled regex patterns against "Header-Name: value" strings. + + Returns the first matching pattern string, or None if nothing matches. + Uses re.compile() which has an internal LRU cache — no per-call overhead + after the first compile. + """ + for pattern in tag_regexes: + for header_str in header_strings: + try: + if re.search(pattern, header_str): + return pattern + except re.error: + verbose_logger.warning( + "tag_regex: invalid pattern %r — skipping", pattern + ) + return None + + def is_valid_deployment_tag( deployment_tags: List[str], request_tags: List[str], match_any: bool = True ) -> bool: @@ -83,30 +107,74 @@ async def get_deployments_for_tag( request_tags = metadata.get("tags") match_any = llm_router_instance.tag_filtering_match_any - new_healthy_deployments = [] - default_deployments = [] - if request_tags: + # Build header strings for regex matching from what the proxy already stores. + # Currently we match against User-Agent; format matches "^User-Agent: claude-code/..." + user_agent = metadata.get("user_agent", "") + header_strings: List[str] = ( + [f"User-Agent: {user_agent}"] if user_agent else [] + ) + + new_healthy_deployments: List[Any] = [] + default_deployments: List[Any] = [] + + has_tag_filter = bool(request_tags) or bool(header_strings) + if has_tag_filter: verbose_logger.debug( - "get_deployments_for_tag routing: router_keys: %s", request_tags + "get_deployments_for_tag routing: request_tags=%s user_agent=%s", + request_tags, + user_agent, ) - # example this can be router_keys=["free", "custom"] for deployment in healthy_deployments: deployment_litellm_params = deployment.get("litellm_params") deployment_tags = deployment_litellm_params.get("tags") + deployment_tag_regex = deployment_litellm_params.get("tag_regex") verbose_logger.debug( - "deployment: %s, deployment_router_keys: %s", - deployment, + "deployment: %s tags: %s tag_regex: %s", + deployment.get("model_name"), deployment_tags, + deployment_tag_regex, ) - if deployment_tags is None: - continue + matched_via: Optional[str] = None + matched_value: Optional[str] = None - if is_valid_deployment_tag(deployment_tags, request_tags, match_any): + # 1. Exact tag match (existing behaviour) + if deployment_tags and request_tags: + if is_valid_deployment_tag(deployment_tags, request_tags, match_any): + matched_via = "tags" + matched_value = next( + (t for t in deployment_tags if t in set(request_tags)), + deployment_tags[0], + ) + + # 2. Regex match against request headers (new) + if matched_via is None and deployment_tag_regex and header_strings: + regex_match = _is_valid_deployment_tag_regex( + deployment_tag_regex, header_strings + ) + if regex_match is not None: + matched_via = "tag_regex" + matched_value = regex_match + + if matched_via is not None: + verbose_logger.debug( + "tag routing match: deployment=%s matched_via=%s matched_value=%s", + deployment.get("model_name"), + matched_via, + matched_value, + ) + # Record provenance in metadata so it flows to SpendLogs + metadata["tag_routing"] = { + "matched_deployment": deployment.get("model_name"), + "matched_via": matched_via, + "matched_value": matched_value, + "request_tags": request_tags or [], + "user_agent": user_agent, + } new_healthy_deployments.append(deployment) - if "default" in deployment_tags: + if deployment_tags and "default" in deployment_tags: default_deployments.append(deployment) if len(new_healthy_deployments) == 0 and len(default_deployments) == 0: diff --git a/litellm/types/router.py b/litellm/types/router.py index f0c1ea5e32a..e8ff2115ff5 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -198,6 +198,11 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): model_info: Optional[Dict] = None mock_response: Optional[Union[str, ModelResponse, Exception, Any]] = None + # tag-based routing + tags: Optional[List[str]] = None + # regex patterns matched against request headers for tag routing + tag_regex: Optional[List[str]] = None + # auto-router params auto_router_config_path: Optional[str] = None auto_router_config: Optional[str] = None @@ -334,6 +339,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): # routing params # use this for tag-based routing tags: Optional[List[str]] + # regex patterns matched against request headers (e.g. "^User-Agent:\\s*claude-code\\/") + tag_regex: Optional[List[str]] # deployment budgets max_budget: Optional[float] diff --git a/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py b/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py new file mode 100644 index 00000000000..cd762dbcdfe --- /dev/null +++ b/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py @@ -0,0 +1,217 @@ +""" +Unit tests for tag_regex routing. + +Tests _is_valid_deployment_tag_regex() and get_deployments_for_tag() with tag_regex +patterns, verifying that regex-based header matching works correctly alongside +existing tag-based routing. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from unittest.mock import MagicMock + +from litellm.router_strategy import tag_based_routing +from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + +_is_valid_deployment_tag_regex = tag_based_routing._is_valid_deployment_tag_regex + + +# --------------------------------------------------------------------------- +# _is_valid_deployment_tag_regex unit tests +# --------------------------------------------------------------------------- + + +def test_regex_matches_claude_code_user_agent(): + """^User-Agent: claude-code/ matches a claude-code UA string.""" + result = _is_valid_deployment_tag_regex( + tag_regexes=[r"^User-Agent: claude-code\/"], + header_strings=["User-Agent: claude-code/1.2.3"], + ) + assert result == r"^User-Agent: claude-code\/" + + +def test_regex_no_match_for_other_ua(): + """Pattern does not match a non-claude-code User-Agent.""" + result = _is_valid_deployment_tag_regex( + tag_regexes=[r"^User-Agent: claude-code\/"], + header_strings=["User-Agent: Mozilla/5.0 (browser)"], + ) + assert result is None + + +def test_regex_returns_first_matching_pattern(): + """When multiple patterns are provided, returns the first match.""" + result = _is_valid_deployment_tag_regex( + tag_regexes=[r"^User-Agent: cursor\/", r"^User-Agent: claude-code\/"], + header_strings=["User-Agent: claude-code/2.0.0"], + ) + assert result == r"^User-Agent: claude-code\/" + + +def test_regex_empty_inputs_return_none(): + """Empty lists return None without errors.""" + assert _is_valid_deployment_tag_regex([], ["User-Agent: claude-code/1.0"]) is None + assert _is_valid_deployment_tag_regex([r"^User-Agent: claude-code\/"], []) is None + + +def test_invalid_regex_skipped_does_not_raise(): + """An invalid regex pattern is skipped (warning logged) — no exception raised.""" + result = _is_valid_deployment_tag_regex( + tag_regexes=["[invalid(regex"], + header_strings=["User-Agent: claude-code/1.0"], + ) + assert result is None + + +def test_regex_matches_version_range(): + """Semver-aware pattern matches multiple versions.""" + pattern = r"^User-Agent: claude-code\/\d" + for ua in ["claude-code/1.0", "claude-code/2.0.0-beta.1", "claude-code/99.0"]: + result = _is_valid_deployment_tag_regex( + tag_regexes=[pattern], + header_strings=[f"User-Agent: {ua}"], + ) + assert result == pattern, f"Expected match for UA: {ua}" + + +# --------------------------------------------------------------------------- +# get_deployments_for_tag integration tests +# --------------------------------------------------------------------------- + +CLAUDE_CODE_DEPLOYMENT = { + "model_name": "claude-sonnet", + "litellm_params": { + "model": "openai/claude-code-deployment", + "api_key": "fake", + "mock_response": "cc", + "tag_regex": [r"^User-Agent: claude-code\/"], + }, + "model_info": {"id": "claude-code-deployment"}, +} + +REGULAR_DEPLOYMENT = { + "model_name": "claude-sonnet", + "litellm_params": { + "model": "openai/regular-deployment", + "api_key": "fake", + "mock_response": "regular", + "tags": ["default"], + }, + "model_info": {"id": "regular-deployment"}, +} + +ALL_DEPLOYMENTS = [CLAUDE_CODE_DEPLOYMENT, REGULAR_DEPLOYMENT] + + +def _make_router_mock(enable_tag_filtering=True, match_any=True): + mock = MagicMock() + mock.enable_tag_filtering = enable_tag_filtering + mock.tag_filtering_match_any = match_any + return mock + + +@pytest.mark.asyncio +async def test_claude_code_ua_routes_to_cc_deployment(): + """claude-code/x.y.z UA → claude-code-deployment via tag_regex.""" + router = _make_router_mock() + result = await get_deployments_for_tag( + llm_router_instance=router, + model="claude-sonnet", + healthy_deployments=ALL_DEPLOYMENTS, + request_kwargs={"metadata": {"user_agent": "claude-code/1.2.3"}}, + ) + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "claude-code-deployment" + + +@pytest.mark.asyncio +async def test_regular_ua_routes_to_default_deployment(): + """Mozilla UA → regular-deployment via default tag fallback.""" + router = _make_router_mock() + result = await get_deployments_for_tag( + llm_router_instance=router, + model="claude-sonnet", + healthy_deployments=ALL_DEPLOYMENTS, + request_kwargs={"metadata": {"user_agent": "Mozilla/5.0 (browser)"}}, + ) + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "regular-deployment" + + +@pytest.mark.asyncio +async def test_no_ua_routes_to_default_deployment(): + """No User-Agent → default deployment.""" + router = _make_router_mock() + result = await get_deployments_for_tag( + llm_router_instance=router, + model="claude-sonnet", + healthy_deployments=ALL_DEPLOYMENTS, + request_kwargs={"metadata": {}}, + ) + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "regular-deployment" + + +@pytest.mark.asyncio +async def test_tag_routing_metadata_written_for_regex_match(): + """tag_routing metadata block is populated when regex matches.""" + router = _make_router_mock() + metadata: dict = {"user_agent": "claude-code/2.0.0-beta.1"} + await get_deployments_for_tag( + llm_router_instance=router, + model="claude-sonnet", + healthy_deployments=ALL_DEPLOYMENTS, + request_kwargs={"metadata": metadata}, + ) + assert "tag_routing" in metadata + tr = metadata["tag_routing"] + assert tr["matched_via"] == "tag_regex" + assert tr["matched_value"] == r"^User-Agent: claude-code\/" + assert tr["user_agent"] == "claude-code/2.0.0-beta.1" + + +@pytest.mark.asyncio +async def test_tag_filtering_disabled_returns_all_deployments(): + """When enable_tag_filtering is False, all deployments returned regardless of UA.""" + router = _make_router_mock(enable_tag_filtering=False) + result = await get_deployments_for_tag( + llm_router_instance=router, + model="claude-sonnet", + healthy_deployments=ALL_DEPLOYMENTS, + request_kwargs={"metadata": {"user_agent": "claude-code/1.0"}}, + ) + assert result == ALL_DEPLOYMENTS + + +@pytest.mark.asyncio +async def test_explicit_tag_match_takes_precedence_over_regex(): + """A deployment with both tags and tag_regex: exact tag match fires first.""" + deployment_with_both = { + "model_name": "claude-sonnet", + "litellm_params": { + "model": "openai/both-deployment", + "api_key": "fake", + "tags": ["premium"], + "tag_regex": [r"^User-Agent: claude-code\/"], + }, + "model_info": {"id": "both-deployment"}, + } + router = _make_router_mock() + metadata: dict = { + "tags": ["premium"], + "user_agent": "claude-code/1.0", + } + result = await get_deployments_for_tag( + llm_router_instance=router, + model="claude-sonnet", + healthy_deployments=[deployment_with_both], + request_kwargs={"metadata": metadata}, + ) + assert len(result) == 1 + tr = metadata.get("tag_routing", {}) + assert tr.get("matched_via") == "tags"