scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

refactor(auth): preserve tool and vector access read order natively

Browse source
This commit is contained in:
Devin AI 2026-11-09 01:36:16 +00:00
parent 4297255514
commit 52075f734f
7 changed files with 468 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

234
litellm-rust/crates/request-auth-core/src/access.rs
View file

@ -76,6 +76,16 @@ pub enum AccessFacts {
selected_trace_present: bool,
other_trace: Option<String>,
},
ToolRequest {
token_present: bool,
call_types: Option<Vec<String>>,
capable_types: Vec<String>,
},
ToolAllowlist {
key_nonempty_list: bool,
team_nonempty_list: Option<bool>,
requested_count: Option<usize>,
},
Tools {
requested: Vec<String>,
key_allowed: Option<Vec<String>>,
@ -90,6 +100,14 @@ pub enum AccessFacts {
session_token: bool,
team_present: bool,
},
VectorRequest {
client_present: bool,
strict: bool,
ids_read: bool,
unreadable: bool,
grants: Option<Vec<Option<Vec<String>>>>,
requested_count: Option<usize>,
},
VectorStores {
requested: Vec<String>,
allowed: Option<Vec<String>>,
@ -115,6 +133,7 @@ pub enum Denial {
Tools,
VectorStores,
MissingVectorGrant,
MalformedVectorIds,
TeamLookup,
}
@ -146,6 +165,17 @@ pub enum AccessRead {
ClientTags,
GuardrailPermission,
UnavailablePolicy,
ToolRoute,
ToolAllowlist,
TeamToolList,
ToolNames,
KeyToolGrant,
TeamToolGrant,
StrictVectorIds,
LegacyVectorIds,
LegacyVectorGrants,
RagVectorIds,
VectorGrants,
}
fn modifies_guardrails(keys: &[String]) -> bool {
@ -173,6 +203,87 @@ fn nonempty(value: &Option<String>) -> bool {
pub fn decide(facts: AccessFacts) -> AccessDecision {
match facts {
AccessFacts::VectorRequest {
client_present,
strict,
ids_read,
unreadable,
grants,
requested_count,
} => {
if !client_present {
return AccessDecision::default();
}
let read = if !ids_read {
Some(if strict {
AccessRead::StrictVectorIds
} else {
AccessRead::LegacyVectorIds
})
} else if unreadable {
let Some(grants) = grants else {
return AccessDecision {
read: Some(AccessRead::LegacyVectorGrants),
..AccessDecision::default()
};
};
return AccessDecision {
denial: grants
.iter()
.flatten()
.any(|grant| !grant.is_empty())
.then_some(Denial::MalformedVectorIds),
..AccessDecision::default()
};
} else {
match requested_count {
None => Some(AccessRead::RagVectorIds),
Some(0) => None,
Some(_) => Some(AccessRead::VectorGrants),
}
};
AccessDecision {
read,
..AccessDecision::default()
}
}
AccessFacts::ToolRequest {
token_present,
call_types,
capable_types,
} => AccessDecision {
read: if !token_present {
None
} else {
match call_types {
None => Some(AccessRead::ToolRoute),
Some(types) => types
.iter()
.any(|kind| capable_types.contains(kind))
.then_some(AccessRead::ToolAllowlist),
}
},
..AccessDecision::default()
},
AccessFacts::ToolAllowlist {
key_nonempty_list,
team_nonempty_list,
requested_count,
} => AccessDecision {
read: if !key_nonempty_list && team_nonempty_list.is_none() {
Some(AccessRead::TeamToolList)
} else if !key_nonempty_list && team_nonempty_list != Some(true) {
None
} else {
match requested_count {
None => Some(AccessRead::ToolNames),
Some(0) => None,
Some(_) if key_nonempty_list => Some(AccessRead::KeyToolGrant),
Some(_) => Some(AccessRead::TeamToolGrant),
}
},
..AccessDecision::default()
},
AccessFacts::ConcurrentFailures { failures, terminal } => AccessDecision {
failure_index: terminal
.iter()
@ -406,6 +517,129 @@ mod tests {
use super::*;
use rstest::rstest;
#[rstest]
#[case::no_database(false, true, false, false, None, None, (None, None))]
#[case::strict_ids(
true,
true,
false,
false,
None,
None,
(Some(AccessRead::StrictVectorIds), None)
)]
#[case::legacy_ids(
true,
false,
false,
false,
None,
None,
(Some(AccessRead::LegacyVectorIds), None)
)]
#[case::unreadable_grants(
true,
false,
true,
true,
None,
None,
(Some(AccessRead::LegacyVectorGrants), None)
)]
#[case::unrestricted_unreadable(true, false, true, true, Some(vec![None, Some(vec![])]), None, (None, None))]
#[case::restricted_unreadable(true, false, true, true, Some(vec![Some(vec!["a".into()])]), None, (None, Some(Denial::MalformedVectorIds)))]
#[case::rag_ids(
true,
true,
true,
false,
None,
None,
(Some(AccessRead::RagVectorIds), None)
)]
#[case::no_ids(true, true, true, false, None, Some(0), (None, None))]
#[case::identity_grants(
true,
true,
true,
false,
None,
Some(2),
(Some(AccessRead::VectorGrants), None)
)]
fn vector_read_order(
#[case] client_present: bool,
#[case] strict: bool,
#[case] ids_read: bool,
#[case] unreadable: bool,
#[case] grants: Option<Vec<Option<Vec<String>>>>,
#[case] requested_count: Option<usize>,
#[case] expected: (Option<AccessRead>, Option<Denial>),
) {
assert_eq!(
decide(AccessFacts::VectorRequest {
client_present,
strict,
ids_read,
unreadable,
grants,
requested_count,
}),
AccessDecision {
read: expected.0,
denial: expected.1,
..AccessDecision::default()
}
);
}
#[rstest]
#[case::no_token(false, None, None)]
#[case::route(true, None, Some(AccessRead::ToolRoute))]
#[case::no_call_type(true, Some(vec![]), None)]
#[case::non_tool_route(true, Some(vec!["embedding".into()]), None)]
#[case::tool_route(true, Some(vec!["embedding".into(), "completion".into()]), Some(AccessRead::ToolAllowlist))]
fn tool_route_reads(
#[case] token_present: bool,
#[case] call_types: Option<Vec<String>>,
#[case] read: Option<AccessRead>,
) {
assert_eq!(
decide(AccessFacts::ToolRequest {
token_present,
call_types,
capable_types: vec!["completion".into()],
})
.read,
read
);
}
#[rstest]
#[case::key_override(true, None, None, Some(AccessRead::ToolNames))]
#[case::team_fallback(false, None, None, Some(AccessRead::TeamToolList))]
#[case::no_allowlist(false, Some(false), None, None)]
#[case::team_enabled(false, Some(true), None, Some(AccessRead::ToolNames))]
#[case::no_tools(true, None, Some(0), None)]
#[case::key_grant(true, None, Some(1), Some(AccessRead::KeyToolGrant))]
#[case::team_grant(false, Some(true), Some(1), Some(AccessRead::TeamToolGrant))]
fn tool_grant_reads(
#[case] key_nonempty_list: bool,
#[case] team_nonempty_list: Option<bool>,
#[case] requested_count: Option<usize>,
#[case] read: Option<AccessRead>,
) {
assert_eq!(
decide(AccessFacts::ToolAllowlist {
key_nonempty_list,
team_nonempty_list,
requested_count,
})
.read,
read
);
}
#[rstest]
#[case::disabled(false, None, None)]
#[case::read_route(true, None, Some(AccessRead::ClientTagRoute))]

50
litellm/_auth/access.py
View file

@ -112,6 +112,20 @@ class TeamTrace:
other_trace: str | None
@dataclass(frozen=True, slots=True)
class ToolRequest:
token_present: bool
call_types: tuple[str, ...] | None = None
capable_types: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class ToolAllowlist:
key_nonempty_list: bool
team_nonempty_list: bool | None = None
requested_count: int | None = None
@dataclass(frozen=True, slots=True)
class Tools:
requested: tuple[str, ...]
@ -130,6 +144,16 @@ class GrantLayers:
team_present: bool
@dataclass(frozen=True, slots=True)
class VectorRequest:
client_present: bool
strict: bool
ids_read: bool = False
unreadable: bool = False
grants: tuple[tuple[str, ...] | None, ...] | None = None
requested_count: int | None = None
@dataclass(frozen=True, slots=True)
class VectorStores:
requested: tuple[str, ...]
@ -160,8 +184,11 @@ AccessFacts: TypeAlias = (
| Guardrails
| AgentTrace
| TeamTrace
| ToolRequest
| ToolAllowlist
| Tools
| GrantLayers
| VectorRequest
| VectorStores
| TeamLookup
)
@ -176,6 +203,7 @@ AccessDenial: TypeAlias = Literal[
"Tools",
"VectorStores",
"MissingVectorGrant",
"MalformedVectorIds",
"TeamLookup",
]
GrantLayer: TypeAlias = Literal["key", "team", "user"]
@ -183,7 +211,27 @@ GrantLayer: TypeAlias = Literal["key", "team", "user"]
@dataclass(frozen=True, slots=True)
class AccessDecision:
read: Literal["UserParam", "ClientTagRoute", "ClientTags", "GuardrailPermission", "UnavailablePolicy"] | None = None
read: (
Literal[
"UserParam",
"ClientTagRoute",
"ClientTags",
"GuardrailPermission",
"UnavailablePolicy",
"ToolRoute",
"ToolAllowlist",
"TeamToolList",
"ToolNames",
"KeyToolGrant",
"TeamToolGrant",
"StrictVectorIds",
"LegacyVectorIds",
"LegacyVectorGrants",
"RagVectorIds",
"VectorGrants",
]
| None
) = None
denial: AccessDenial | None = None
rejected: tuple[str, ...] = ()
layers: tuple[GrantLayer, ...] = ()

60
litellm/proxy/auth/auth_checks.py
View file

@ -27,7 +27,7 @@ from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack
import litellm
from litellm._auth import access as access_policy
from litellm._auth import common_flow
from litellm._auth.failure import NativeAuthFailure
from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict
@ -782,21 +782,36 @@ async def check_tools_allowlist(
get_call_types_for_route,
)
if valid_token is None:
if access_policy.decide(access_policy.ToolRequest(valid_token is not None)).read is None:
return
assert valid_token is not None
call_types: Final = get_call_types_for_route(route)
if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types):
route_plan: Final = access_policy.decide(
access_policy.ToolRequest(True, tuple(ct.value for ct in call_types or ()), tuple(TOOL_CAPABLE_CALL_TYPES))
)
if route_plan.read is None:
return
key_meta: Final = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {}
team_meta: Final = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {}
key_allowed: Final = key_meta.get("allowed_tools")
team_allowed: Final = team_meta.get("allowed_tools")
effective: Final = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed
if not isinstance(effective, list) or len(effective) == 0:
key_facts: Final = access_policy.ToolAllowlist(isinstance(key_allowed, list) and len(key_allowed) > 0)
key_plan: Final = access_policy.decide(key_facts)
facts: Final = (
replace(key_facts, team_nonempty_list=isinstance(team_allowed, list) and len(team_allowed) > 0)
if key_plan.read == "TeamToolList"
else key_facts
)
if access_policy.decide(facts).read is None:
return
tool_names: Final = extract_request_tool_names(route, request_body)
if not tool_names:
grant_plan: Final = access_policy.decide(replace(facts, requested_count=len(tool_names)))
if grant_plan.read is None:
return
if grant_plan.read not in ("KeyToolGrant", "TeamToolGrant"):
raise NativeAuthRuntimeError("Unexpected tool-grant decision")
effective: Final = key_allowed if grant_plan.read == "KeyToolGrant" else team_allowed
assert isinstance(effective, list)
decision: Final = access_policy.decide(
access_policy.Tools(
requested=tuple(tool_names),
@ -4610,7 +4625,7 @@ async def can_key_call_model(
async def can_key_call_resolved_model(
model: str,
model: str | Sequence[str],
llm_model_list: Sequence[object] | None,
valid_token: UserAPIKeyAuth,
llm_router: RegistryReader | None,
@ -5824,7 +5839,21 @@ def _legacy_requested_vector_store_ids(
def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, LiteLLM_ObjectPermissionTable | None]]) -> bool:
return any(grant is not None and bool(grant.vector_stores) for _, grant in grants)
return (
access_policy.decide(
access_policy.VectorRequest(
client_present=True,
strict=False,
ids_read=True,
unreadable=True,
grants=tuple(
tuple(grant.vector_stores) if grant is not None and grant.vector_stores is not None else None
for _, grant in grants
),
)
).denial
is not None
)
def _request_drop_policies(
@ -5932,28 +5961,35 @@ async def vector_store_access_check(
#########################################################
# Get the vector store the user is trying to access
#########################################################
if prisma_client is None:
facts: Final = access_policy.VectorRequest(prisma_client is not None, deny_by_default)
initial: Final = access_policy.decide(facts)
if initial.read is None:
verbose_proxy_logger.debug("Prisma client not found, skipping vector store access check")
return True
assert prisma_client is not None
typed_request_body: Final = _typed_request_body(request_body)
requested_ids: Final = (
_strict_requested_vector_store_ids(typed_request_body)
if deny_by_default
if initial.read == "StrictVectorIds"
else _legacy_requested_vector_store_ids(
typed_request_body, lambda: _request_drop_policies(typed_request_body, llm_router, valid_token)
)
)
if isinstance(requested_ids, _UnreadableVectorStoreIds):
ids_plan: Final = access_policy.decide(
replace(facts, ids_read=True, unreadable=isinstance(requested_ids, _UnreadableVectorStoreIds))
)
if ids_plan.read == "LegacyVectorGrants":
if _restricts_vector_stores(
await _identity_grants(valid_token, team_object, None, prisma_client, user_api_key_cache)
):
raise _malformed_vector_store_ids()
return True
assert not isinstance(requested_ids, _UnreadableVectorStoreIds)
rag_vector_store_id: Final = _get_rag_query_vector_store_id(typed_request_body)
rag_ids: Final = (rag_vector_store_id,) if rag_vector_store_id is not None else ()
vector_store_ids_to_run: Final = tuple(dict.fromkeys((*requested_ids, *rag_ids)))
if not vector_store_ids_to_run:
if access_policy.decide(replace(facts, ids_read=True, requested_count=len(vector_store_ids_to_run))).read is None:
return True
#########################################################

6
litellm/proxy/auth/common_budget_lookups.py
View file

@ -22,9 +22,9 @@ from litellm.proxy.auth.common_budget_adapter import (
resolve_effective_team_member_budget,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.redis.stores.spend_counters import model_access_group_counter
from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import model_access_group_counter
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
@ -200,9 +200,7 @@ async def read_access_group_budget_row(name: str, row: ModelAccessGroupBudget |
limit: Final = row.max_budget if row is not None else None
facts, _decision = await read_budget(
BudgetFacts("AccessGroup", limit),
CounterRead(
model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True
),
CounterRead(model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True),
)
return AccessGroupBudgetReading(name, facts)

11
litellm/proxy/auth/resolved_model_adapter.py
View file

@ -44,7 +44,7 @@ class ResolvedModelReading:
@dataclass(frozen=True, slots=True)
class _Source:
model: str
model: str | Sequence[str]
token: UserAPIKeyAuth
router: RegistryReader | None
prisma: PrismaClient | None
@ -57,7 +57,7 @@ class _Source:
async def read_resolved_model_facts(
model: str,
model: str | Sequence[str],
valid_token: UserAPIKeyAuth,
llm_router: RegistryReader | None,
prisma_client: PrismaClient | None,
@ -65,7 +65,7 @@ async def read_resolved_model_facts(
proxy_logging_obj: ProxyLogging,
) -> ResolvedModelReading:
facts: Final = ResolvedModelFacts(
model=model,
model=str(model),
config_nonempty=bool(_Config.model_validate(valid_token, from_attributes=True).config),
key_models=object_model_grants(valid_token),
team_id=valid_token.team_id,
@ -83,7 +83,10 @@ async def read_resolved_model_facts(
async def can_key_call_resolved_model(
model: str, llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None
model: str | Sequence[str],
llm_model_list: Sequence[object] | None,
valid_token: UserAPIKeyAuth,
llm_router: RegistryReader | None,
) -> None:
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache

73
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -2320,6 +2320,18 @@ async def test_every_requested_vector_store_needs_a_grant(
)
@pytest.mark.asyncio
@pytest.mark.parametrize("deny_by_default", (False, True), ids=("legacy", "strict"))
async def test_missing_database_skips_malformed_vector_request(deny_by_default: bool) -> None:
with patch("litellm.proxy.proxy_server.prisma_client", None):
assert await vector_store_access_check(
request_body={"vector_store_ids": {"malformed": "ids"}},
valid_token=_virtual_key(),
team_object=None,
deny_by_default=deny_by_default,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"valid_token",
@ -2432,6 +2444,67 @@ async def test_file_search_inside_additional_tools_input_item_needs_a_vector_sto
)
@pytest.mark.parametrize("strict", [False, True], ids=["legacy", "strict"])
@pytest.mark.parametrize(
("input_items", "expected"),
[
(None, ("top", "shared", "tool", "shared")),
("not-a-list", ("top", "shared", "tool", "shared")),
([{"type": "additional_tools", "tools": 1}], ("top", "shared", "tool", "shared")),
(
[
{"type": "message", "tools": [{"vector_store_ids": ["ignored"]}]},
{"type": "additional_tools", "tools": [{"vector_store_ids": ["input", "shared"]}]},
{"type": "additional_tools", "tools": [{"vector_store_ids": ["last"]}]},
],
("top", "shared", "tool", "shared", "input", "shared", "last"),
),
],
ids=["no-input", "string-input", "malformed-tools", "ordered-duplicates"],
)
def test_vector_store_raw_reads_keep_top_level_tools_then_additional_tools(
strict: bool, input_items: object, expected: tuple[str, ...]
) -> None:
from litellm.proxy.auth.auth_checks import _legacy_requested_vector_store_ids, _strict_requested_vector_store_ids
body: Final = {
"vector_store_ids": ["top", "shared"],
"tools": [{"vector_store_ids": ["tool", "shared"]}],
"input": input_items,
}
result: Final = (
_strict_requested_vector_store_ids(body) if strict else _legacy_requested_vector_store_ids(body, lambda: ())
)
assert result == expected
@pytest.mark.parametrize("strict", [False, True], ids=["legacy", "strict"])
def test_malformed_top_level_vector_ids_do_not_read_additional_tools(strict: bool) -> None:
from litellm.proxy.auth.auth_checks import (
_legacy_requested_vector_store_ids,
_strict_requested_vector_store_ids,
_UnreadableVectorStoreIds,
)
with patch(
"litellm.proxy.auth.auth_checks.additional_tools_in_input",
side_effect=AssertionError("unreadable top-level IDs read input tools"),
) as input_reader:
if strict:
with pytest.raises(ProxyException) as caught:
_strict_requested_vector_store_ids({"vector_store_ids": 42})
assert (caught.value.code, caught.value.type, caught.value.param) == (
"400",
"invalid_request_error",
"vector_store_ids",
)
else:
assert isinstance(
_legacy_requested_vector_store_ids({"vector_store_ids": 42}, lambda: ()), _UnreadableVectorStoreIds
)
input_reader.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"request_body",

63
tests/unit/proxy/test_tools_allowlist_enforcement.py
View file

@ -6,14 +6,14 @@ Covers:
- extract_request_tool_names: OpenAI chat, responses, Anthropic, generate_content, MCP
"""
from unittest.mock import AsyncMock, MagicMock, patch
from typing import Final
from unittest.mock import MagicMock, patch
import pytest
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import check_tools_allowlist
from litellm.proxy.guardrails.tool_name_extraction import (
TOOL_CAPABLE_CALL_TYPES,
extract_request_tool_names,
)
@ -58,9 +58,7 @@ class TestExtractRequestToolNames:
{"type": "function", "name": "get_current_weather", "description": "x"},
]
}
assert extract_request_tool_names("/v1/responses", data) == [
"get_current_weather"
]
assert extract_request_tool_names("/v1/responses", data) == ["get_current_weather"]
def test_openai_responses_mcp_tools(self):
data = {
@ -129,9 +127,7 @@ class TestExtractRequestToolNames:
},
]
}
assert extract_request_tool_names("/generate_content", data) == [
"schedule_meeting"
]
assert extract_request_tool_names("/generate_content", data) == ["schedule_meeting"]
def test_mcp_call_tool_name(self):
data = {"name": "my_tool", "arguments": {}}
@ -164,6 +160,57 @@ class TestExtractRequestToolNames:
class TestCheckToolsAllowlist:
"""Test allowlist enforcement in auth (no DB in hot path)."""
@pytest.mark.asyncio
@pytest.mark.parametrize("key", [None, [], "forbidden", {"tool": "forbidden"}])
@pytest.mark.parametrize("team", [None, [], "forbidden", {"tool": "forbidden"}])
async def test_disabled_allowlists_never_extract_tool_names(self, key: object, team: object) -> None:
with patch(
"litellm.proxy.auth.auth_checks.extract_request_tool_names",
side_effect=AssertionError("disabled allowlist read request tools"),
) as extractor:
await check_tools_allowlist(
request_body={"tools": object()},
valid_token=_token(metadata={"allowed_tools": key}, team_metadata={"allowed_tools": team}),
team_object=None,
route="/v1/responses",
)
extractor.assert_not_called()
@pytest.mark.asyncio
async def test_empty_tool_names_do_not_convert_allowlist_members(self) -> None:
member: Final = MagicMock()
with patch("litellm.proxy.auth.auth_checks.extract_request_tool_names", return_value=[]) as extractor:
await check_tools_allowlist(
request_body={},
valid_token=_token(metadata={"allowed_tools": [member]}),
team_object=None,
route="/v1/responses",
)
assert extractor.call_count == 1
member.__str__.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize("key", [None, [], "ignored", {"tool": "ignored"}])
async def test_team_members_are_stringified_only_after_tool_names(self, key: object) -> None:
member: Final = MagicMock()
member.__str__.return_value = "allowed"
def names(route: str, body: object) -> list[str]:
member.__str__.assert_not_called()
return ["allowed", "123", "None", "False"]
with patch("litellm.proxy.auth.auth_checks.extract_request_tool_names", side_effect=names):
await check_tools_allowlist(
request_body={},
valid_token=_token(
metadata={"allowed_tools": key},
team_metadata={"allowed_tools": [member, 123, None, False]},
),
team_object=None,
route="/v1/responses",
)
member.__str__.assert_called_once_with()
@pytest.mark.asyncio
async def test_no_allowlist_passes(self):
token = _token(metadata={}, team_metadata={})