refactor: stage personal model listings and agent ceilings in native policy

litellm-rust/crates/python-bridge/src/lib.rs

@@ -56,6 +56,9 @@ mod _native {
         request_auth_personal_model_request,
         request_auth_member_models,
         request_auth_resolved_models,
+        request_auth_personal_model_listing,
+        request_auth_agent_ceiling,
+        request_auth_agent_group_load,
     };
     #[pymodule_export]
     use crate::router::{RouterCall, RouterCoreHandle};

litellm-rust/crates/python-bridge/src/request_auth.rs

@@ -287,3 +287,23 @@ pub(crate) fn request_auth_resolved_models(py: Python<'_>, facts: &Bound<'_, PyA
     let decision = py.detach(move || core::resolved_models::decide(&facts)).map_err(map_error)?;
     to_py(py, &decision)
 }
+
+#[pyfunction]
+pub(crate) fn request_auth_personal_model_listing(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+    let facts: core::personal_model_listing::Facts = from_py(facts)?;
+    let decision = py.detach(move || core::personal_model_listing::decide(&facts)).map_err(map_error)?;
+    to_py(py, &decision)
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_agent_ceiling(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+    let facts: core::agent_ceilings::Facts = from_py(facts)?;
+    let decision = py.detach(move || core::agent_ceilings::decide(&facts)).map_err(map_error)?;
+    to_py(py, &decision)
+}
+
+#[pyfunction]
+pub(crate) fn request_auth_agent_group_load(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+    let facts: core::agent_ceilings::LoadFacts = from_py(facts)?;
+    to_py(py, &py.detach(move || core::agent_ceilings::load(&facts)))
+}

litellm-rust/crates/request-auth-core/src/agent_ceilings.rs

@@ -0,0 +1,151 @@
+use crate::Error;
+use serde::{Deserialize, Serialize};
+use std::collections::BTreeSet;
+
+#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)]
+pub enum Status {
+    Found,
+    Missing,
+    Unavailable,
+}
+
+#[derive(Clone, Deserialize)]
+pub struct Group {
+    pub status: Status,
+    pub models: Vec<String>,
+    pub mcp_server_ids: Vec<String>,
+    pub agent_ids: Vec<String>,
+}
+
+#[derive(Clone, Deserialize)]
+pub struct Facts {
+    pub ids: Option<Vec<String>>,
+    pub groups: Option<Vec<Group>>,
+    pub check_db_only: bool,
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize)]
+#[serde(tag = "kind")]
+pub enum Decision {
+    ReadIds,
+    ReadGroups {
+        authoritative: bool,
+    },
+    Complete {
+        present: bool,
+        ids: Vec<String>,
+        models: BTreeSet<String>,
+        mcp_server_ids: BTreeSet<String>,
+        agent_ids: BTreeSet<String>,
+        indices: Vec<usize>,
+    },
+}
+
+pub fn decide(facts: &Facts) -> Result<Decision, Error> {
+    let Some(ids) = &facts.ids else {
+        return Ok(Decision::ReadIds);
+    };
+    let empty = |present| Decision::Complete {
+        present,
+        ids: vec![],
+        models: BTreeSet::new(),
+        mcp_server_ids: BTreeSet::new(),
+        agent_ids: BTreeSet::new(),
+        indices: vec![],
+    };
+    if ids.is_empty() {
+        return Ok(empty(false));
+    }
+    let Some(groups) = &facts.groups else {
+        return Ok(Decision::ReadGroups {
+            authoritative: facts.check_db_only,
+        });
+    };
+    if groups.len() != ids.len() {
+        return Err(Error::InvalidAnswer);
+    }
+    let indices: Vec<_> = groups
+        .iter()
+        .enumerate()
+        .filter_map(|(index, group)| (group.status == Status::Found).then_some(index))
+        .collect();
+    Ok(Decision::Complete {
+        present: true,
+        ids: ids
+            .iter()
+            .zip(groups)
+            .filter(|(_, group)| group.status != Status::Missing)
+            .map(|(id, _)| id.clone())
+            .collect(),
+        models: indices
+            .iter()
+            .flat_map(|index| groups[*index].models.iter().cloned())
+            .collect(),
+        mcp_server_ids: indices
+            .iter()
+            .flat_map(|index| groups[*index].mcp_server_ids.iter().cloned())
+            .collect(),
+        agent_ids: indices
+            .iter()
+            .flat_map(|index| groups[*index].agent_ids.iter().cloned())
+            .collect(),
+        indices,
+    })
+}
+
+#[derive(Clone, Copy, Deserialize, strum::VariantArray)]
+pub enum ReadStatus {
+    Success,
+    NotFound,
+    HttpError,
+}
+
+#[derive(Clone, Deserialize)]
+pub struct LoadFacts {
+    pub database_present: bool,
+    pub check_db_only: bool,
+    pub status: Option<ReadStatus>,
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize)]
+#[serde(tag = "kind")]
+pub enum LoadDecision {
+    Read,
+    Return,
+    DatabaseUnavailable,
+    HttpUnavailable,
+    Raise,
+}
+
+pub fn load(facts: &LoadFacts) -> LoadDecision {
+    if !facts.database_present {
+        return LoadDecision::DatabaseUnavailable;
+    }
+    match facts.status {
+        None => LoadDecision::Read,
+        Some(ReadStatus::Success) => LoadDecision::Return,
+        Some(ReadStatus::NotFound) => LoadDecision::Raise,
+        Some(ReadStatus::HttpError) if facts.check_db_only => LoadDecision::Raise,
+        Some(ReadStatus::HttpError) => LoadDecision::HttpUnavailable,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::soft(false, LoadDecision::HttpUnavailable)]
+    #[case::authoritative(true, LoadDecision::Raise)]
+    fn http_lookup_policy(#[case] check_db_only: bool, #[case] expected: LoadDecision) {
+        assert_eq!(
+            load(&LoadFacts {
+                database_present: true,
+                check_db_only,
+                status: Some(ReadStatus::HttpError)
+            }),
+            expected
+        );
+    }
+}

litellm-rust/crates/request-auth-core/src/lib.rs

@@ -23,6 +23,8 @@ pub mod personal_owner;
 pub mod personal_model_request;
 pub mod member_models;
 pub mod resolved_models;
+pub mod personal_model_listing;
+pub mod agent_ceilings;
 pub mod owner_teams;
 pub mod protocol;
 pub mod route_access;

litellm-rust/crates/request-auth-core/src/personal_model_listing.rs

@@ -0,0 +1,81 @@
+use crate::{Error, personal_model_request::Loaded, personal_models, personal_owner};
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Deserialize)]
+pub struct Facts {
+    pub owner: personal_owner::Facts,
+    pub names: Option<Vec<String>>,
+    pub models: Vec<Loaded>,
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize)]
+#[serde(tag = "kind")]
+pub enum Decision {
+    ReadOwner,
+    ReadNames,
+    ReadModel { index: usize, teams: Vec<usize> },
+    Complete { selected: Vec<usize> },
+}
+
+pub fn decide(facts: &Facts) -> Result<Decision, Error> {
+    let teams = match personal_owner::decide(&facts.owner)? {
+        personal_owner::Decision::Complete { selected, .. } => selected,
+        _ => return Ok(Decision::ReadOwner),
+    };
+    if facts
+        .owner
+        .teams
+        .as_ref()
+        .is_none_or(|teams| teams.loads.is_empty())
+    {
+        return Ok(Decision::Complete { selected: vec![] });
+    }
+    let Some(names) = &facts.names else {
+        return Ok(Decision::ReadNames);
+    };
+    let mut selected = vec![];
+    for index in 0..names.len() {
+        let Some(loaded) = facts.models.get(index) else {
+            return Ok(Decision::ReadModel { index, teams });
+        };
+        match personal_models::decide(personal_models::Request::Loaded {
+            present: loaded.present.clone(),
+            scopes: loaded.scopes.clone(),
+        })? {
+            personal_models::Decision::Complete { allowed } => {
+                if allowed {
+                    selected.push(index);
+                }
+            }
+            _ => return Err(Error::InvalidAnswer),
+        }
+    }
+    Ok(Decision::Complete { selected })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::missing_database(false, Decision::Complete { selected: vec![] })]
+    #[case::load(true, Decision::ReadOwner)]
+    fn owner_before_names(#[case] database_present: bool, #[case] expected: Decision) {
+        let facts = Facts {
+            owner: personal_owner::Facts {
+                user_present: true,
+                database_present,
+                cache_present: true,
+                logging_present: true,
+                owner_loaded: false,
+                owner_present: false,
+                owner_failed: false,
+                teams: None,
+            },
+            names: None,
+            models: vec![],
+        };
+        assert_eq!(decide(&facts).unwrap(), expected);
+    }
+}

litellm/_auth/agent_ceilings.py

@@ -0,0 +1,60 @@
+from dataclasses import asdict, dataclass
+from typing import Final, Literal
+
+from pydantic import TypeAdapter
+
+from litellm._auth.failure import native_boundary
+from litellm.rust_bridge import _native
+
+
+@dataclass(frozen=True, slots=True)
+class AgentCeilingGroup:
+    status: Literal["Found", "Missing", "Unavailable"]
+    models: tuple[str, ...] = ()
+    mcp_server_ids: tuple[str, ...] = ()
+    agent_ids: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class AgentCeilingFacts:
+    ids: tuple[str, ...] | None = None
+    groups: tuple[AgentCeilingGroup, ...] | None = None
+    check_db_only: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class AgentCeilingDecision:
+    kind: Literal["ReadIds", "ReadGroups", "Complete"]
+    authoritative: bool = False
+    present: bool = False
+    ids: tuple[str, ...] = ()
+    models: frozenset[str] = frozenset()
+    mcp_server_ids: frozenset[str] = frozenset()
+    agent_ids: frozenset[str] = frozenset()
+    indices: tuple[int, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class AgentGroupLoadFacts:
+    database_present: bool
+    check_db_only: bool = False
+    status: Literal["Success", "NotFound", "HttpError"] | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class AgentGroupLoadDecision:
+    kind: Literal["Read", "Return", "DatabaseUnavailable", "HttpUnavailable", "Raise"]
+
+
+_DECISION: Final = TypeAdapter(AgentCeilingDecision)
+_LOAD: Final = TypeAdapter(AgentGroupLoadDecision)
+
+
+@native_boundary
+def agent_ceiling(facts: AgentCeilingFacts) -> AgentCeilingDecision:
+    return _DECISION.validate_python(_native.request_auth_agent_ceiling(asdict(facts)))
+
+
+@native_boundary
+def agent_group_load(facts: AgentGroupLoadFacts) -> AgentGroupLoadDecision:
+    return _LOAD.validate_python(_native.request_auth_agent_group_load(asdict(facts)))

litellm/_auth/personal_model_listing.py

@@ -0,0 +1,32 @@
+from dataclasses import asdict, dataclass
+from typing import Final, Literal
+
+from pydantic import TypeAdapter
+
+from litellm._auth.failure import native_boundary
+from litellm._auth.personal_models import LoadedTeamModelFacts
+from litellm._auth.personal_owner import PersonalOwnerFacts
+from litellm.rust_bridge import _native
+
+
+@dataclass(frozen=True, slots=True)
+class PersonalModelListingFacts:
+    owner: PersonalOwnerFacts
+    names: tuple[str, ...] | None = None
+    models: tuple[LoadedTeamModelFacts, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class PersonalModelListingDecision:
+    kind: Literal["ReadOwner", "ReadNames", "ReadModel", "Complete"]
+    index: int = 0
+    teams: tuple[int, ...] = ()
+    selected: tuple[int, ...] = ()
+
+
+_DECISION: Final = TypeAdapter(PersonalModelListingDecision)
+
+
+@native_boundary
+def personal_model_listing(facts: PersonalModelListingFacts) -> PersonalModelListingDecision:
+    return _DECISION.validate_python(_native.request_auth_personal_model_listing(asdict(facts)))

litellm/proxy/agent_endpoints/auth/agent_access_groups.py

@@ -1,11 +1,7 @@
-import asyncio
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Final, TypeAlias
 
-from fastapi import HTTPException
-
-from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import LiteLLM_AccessGroupTable
 
 if TYPE_CHECKING:
@@ -31,13 +27,6 @@ class AgentAccessGroupCeiling:
 CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]]
 
 
-class _MissingGroup:
-    __slots__ = ()
-
-
-_MISSING: Final = _MissingGroup()
-
-
 async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
     from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
 
@@ -46,29 +35,9 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
 
 
 async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup:
-    from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object
-    from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
+    from litellm.proxy.auth.agent_ceiling_adapter import load_agent_access_group
 
-    if prisma_client is None:
-        verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id)
-        return None
-    try:
-        return await get_access_object(
-            access_group_id=access_group_id,
-            prisma_client=prisma_client,
-            user_api_key_cache=user_api_key_cache,
-            proxy_logging_obj=proxy_logging_obj,
-            check_db_only=check_db_only,
-        )
-    except AccessGroupNotFoundError:
-        raise
-    except HTTPException as e:
-        if check_db_only:
-            raise
-        verbose_proxy_logger.warning(
-            "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail
-        )
-        return None
+    return await load_agent_access_group(access_group_id, check_db_only=check_db_only)
 
 
 async def resolve_agent_access_group_ceiling(
@@ -82,31 +51,10 @@ async def resolve_agent_access_group_ceiling(
     provably does not exist is dropped from the ceiling, it grants and bounds nothing, while a group
     that exists but could not be read stays named and loads as ``None``. The ceiling itself stays once
     any group is named, so an agent whose every group is missing is still capped to nothing."""
-    named_ids: Final = await load_access_group_ids(agent_id)
-    if not named_ids:
-        return None
+    from litellm.proxy.auth.agent_ceiling_adapter import project_agent_ceiling, read_agent_ceiling_facts
 
-    from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
-
-    async def load(group_id: str) -> LoadedAccessGroup | _MissingGroup:
-        try:
-            return await (
-                _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id)
-            )
-        except AccessGroupNotFoundError:
-            return _MISSING
-
-    results: Final = await asyncio.gather(*(load(group_id) for group_id in named_ids))
-    access_group_ids: Final = tuple(
-        group_id for group_id, result in zip(named_ids, results, strict=True) if not isinstance(result, _MissingGroup)
-    )
-    groups: Final = tuple(result for result in results if isinstance(result, LiteLLM_AccessGroupTable))
-    return AgentAccessGroupCeiling(
-        access_group_ids=access_group_ids,
-        models=frozenset(model for group in groups for model in group.access_model_names),
-        mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids),
-        agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids),
-        access_groups=groups,
+    return project_agent_ceiling(
+        await read_agent_ceiling_facts(agent_id, load_access_group_ids, load_access_group, check_db_only=check_db_only)
     )
 
 

litellm/proxy/auth/agent_ceiling_adapter.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Final
+
+from fastapi import HTTPException
+
+from litellm._auth.agent_ceilings import (
+    AgentCeilingFacts,
+    AgentCeilingGroup,
+    AgentGroupLoadFacts,
+    agent_ceiling,
+    agent_group_load,
+)
+from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import LiteLLM_AccessGroupTable
+
+if TYPE_CHECKING:
+    from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
+        AccessGroupIdsLoader,
+        AccessGroupLoader,
+        AgentAccessGroupCeiling,
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class AgentCeilingReading:
+    facts: AgentCeilingFacts
+    rows: tuple[LiteLLM_AccessGroupTable | None, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class _Source:
+    agent_id: str
+    ids: AccessGroupIdsLoader
+    group: AccessGroupLoader
+
+
+async def read_agent_ceiling_facts(
+    agent_id: str,
+    load_access_group_ids: AccessGroupIdsLoader,
+    load_access_group: AccessGroupLoader,
+    *,
+    check_db_only: bool = False,
+) -> AgentCeilingReading:
+    return await _read(
+        AgentCeilingReading(AgentCeilingFacts(check_db_only=check_db_only)),
+        _Source(agent_id, load_access_group_ids, load_access_group),
+    )
+
+
+def project_agent_ceiling(reading: AgentCeilingReading) -> AgentAccessGroupCeiling | None:
+    from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling
+
+    decision: Final = agent_ceiling(reading.facts)
+    if decision.kind != "Complete":
+        raise NativeAuthRuntimeError("incomplete agent ceiling facts")
+    if not decision.present:
+        return None
+    return AgentAccessGroupCeiling(
+        decision.ids,
+        decision.models,
+        decision.mcp_server_ids,
+        decision.agent_ids,
+        tuple(_row(reading, index) for index in decision.indices),
+    )
+
+
+def _row(reading: AgentCeilingReading, index: int) -> LiteLLM_AccessGroupTable:
+    row: Final = reading.rows[index]
+    if row is None:
+        raise NativeAuthRuntimeError("agent ceiling selected a missing row")
+    return row
+
+
+async def _group(group_id: str, loader: AccessGroupLoader) -> tuple[AgentCeilingGroup, LiteLLM_AccessGroupTable | None]:
+    from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
+
+    try:
+        row: Final = await loader(group_id)
+    except NativeAuthFailure:
+        raise
+    except AccessGroupNotFoundError:
+        return AgentCeilingGroup("Missing"), None
+    if not isinstance(row, LiteLLM_AccessGroupTable):
+        return AgentCeilingGroup("Unavailable"), None
+    return AgentCeilingGroup(
+        "Found", tuple(row.access_model_names), tuple(row.access_mcp_server_ids), tuple(row.access_agent_ids)
+    ), row
+
+
+async def _authoritative(group_id: str) -> LiteLLM_AccessGroupTable | None:
+    from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
+        _load_access_group,  # pyright: ignore[reportPrivateUsage]  # authoritative raw reader preserves the public injected-loader contract
+    )
+
+    return await _load_access_group(group_id, check_db_only=True)
+
+
+async def _read(reading: AgentCeilingReading, source: _Source) -> AgentCeilingReading:
+    decision: Final = agent_ceiling(reading.facts)
+    match decision.kind:
+        case "Complete":
+            return reading
+        case "ReadIds":
+            ids: Final = await source.ids(source.agent_id)
+            return await _read(replace(reading, facts=replace(reading.facts, ids=ids)), source)
+        case "ReadGroups":
+            if reading.facts.ids is None:
+                raise NativeAuthRuntimeError("agent group reads without named IDs")
+            loader: Final = _authoritative if decision.authoritative else source.group
+            rows: Final = tuple(await asyncio.gather(*(_group(group_id, loader) for group_id in reading.facts.ids)))
+            return AgentCeilingReading(
+                replace(reading.facts, groups=tuple(row[0] for row in rows)), tuple(row[1] for row in rows)
+            )
+
+
+async def load_agent_access_group(
+    access_group_id: str, *, check_db_only: bool = False
+) -> LiteLLM_AccessGroupTable | None:
+    from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object
+    from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
+
+    initial: Final = AgentGroupLoadFacts(prisma_client is not None, check_db_only)
+    initial_decision: Final = agent_group_load(initial)
+    if initial_decision.kind == "DatabaseUnavailable":
+        verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id)
+        return None
+    if initial_decision.kind != "Read":
+        raise NativeAuthRuntimeError("invalid initial agent group read decision")
+    try:
+        return await get_access_object(
+            access_group_id=access_group_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+            check_db_only=check_db_only,
+        )
+    except NativeAuthFailure:
+        raise
+    except AccessGroupNotFoundError:
+        if agent_group_load(replace(initial, status="NotFound")).kind != "Raise":
+            raise NativeAuthRuntimeError("invalid missing agent group decision") from None
+        raise
+    except HTTPException as error:
+        decision: Final = agent_group_load(replace(initial, status="HttpError"))
+        if decision.kind == "Raise":
+            raise
+        if decision.kind != "HttpUnavailable":
+            raise NativeAuthRuntimeError("invalid failed agent group decision") from error
+        verbose_proxy_logger.warning(
+            "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, error.detail
+        )
+        return None

litellm/proxy/auth/model_access_adapter.py

@@ -6,7 +6,7 @@ from fastapi import HTTPException
 from pydantic import BaseModel
 
 import litellm
-from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError
+from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary
 from litellm._auth.model_access import (
     ModelAccessDecision,
     ModelAccessFacts,
@@ -45,6 +45,7 @@ class ModelLists(BaseModel, frozen=True):
     models: tuple[str, ...] = ()
 
 
+@native_boundary
 def object_model_grants(source: object) -> tuple[str, ...]:
     return ModelLists.model_validate(source, from_attributes=True).models
 
@@ -71,6 +72,7 @@ def _record(model: str, router: RegistryReader | None, team_id: str | None) -> M
     )
 
 
+@native_boundary
 def project_model_facts(
     model: str | Sequence[str],
     llm_router: RegistryReader | None,

litellm/proxy/auth/personal_model_listing_adapter.py

@@ -0,0 +1,111 @@
+from collections.abc import Sequence
+from dataclasses import dataclass, replace
+from typing import Final
+
+from litellm._auth.failure import NativeAuthRuntimeError
+from litellm._auth.personal_model_listing import PersonalModelListingFacts, personal_model_listing
+from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner
+from litellm.proxy._types import LiteLLM_TeamTableCachedObj, UserAPIKeyAuth
+from litellm.proxy.auth.personal_model_adapter import read_loaded_team_model_facts
+from litellm.proxy.auth.personal_owner_adapter import read_personal_owner_facts, selected_team
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+from litellm.router_api import RegistryReader
+
+
+@dataclass(frozen=True, slots=True)
+class _Source:
+    names: Sequence[str]
+    token: UserAPIKeyAuth
+    router: RegistryReader | None
+    prisma: PrismaClient | None
+    cache: UserApiKeyCache | None
+    logging: ProxyLogging | None
+    teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = ()
+
+
+async def personal_key_listing_teams(
+    user_api_key_dict: UserAPIKeyAuth,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache | None,
+    proxy_logging_obj: ProxyLogging | None,
+    *,
+    check_db_only: bool = False,
+) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]:
+    reading: Final = await read_personal_owner_facts(
+        user_api_key_dict,
+        prisma_client,
+        user_api_key_cache,
+        proxy_logging_obj,
+        check_db_only=check_db_only,
+        purpose="Models",
+    )
+    decision: Final = personal_owner(reading.facts)
+    if decision.kind != "Complete":
+        raise NativeAuthRuntimeError("incomplete listing team facts")
+    return tuple(team if index in decision.selected else None for index, team in enumerate(reading.teams))
+
+
+async def read_personal_model_listing_facts(
+    models: Sequence[str],
+    user_api_key_dict: UserAPIKeyAuth,
+    llm_router: RegistryReader | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache | None,
+    proxy_logging_obj: ProxyLogging | None,
+) -> PersonalModelListingFacts:
+    owner: Final = PersonalOwnerFacts(
+        user_api_key_dict.user_id is not None,
+        prisma_client is not None,
+        user_api_key_cache is not None,
+        proxy_logging_obj is not None,
+    )
+    return await _read(
+        PersonalModelListingFacts(owner),
+        _Source(models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj),
+    )
+
+
+async def personal_key_team_visible_models(
+    models: Sequence[str],
+    user_api_key_dict: UserAPIKeyAuth,
+    llm_router: RegistryReader | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache | None,
+    proxy_logging_obj: ProxyLogging | None,
+) -> tuple[str, ...]:
+    facts: Final = await read_personal_model_listing_facts(
+        models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj
+    )
+    decision: Final = personal_model_listing(facts)
+    if decision.kind != "Complete":
+        raise NativeAuthRuntimeError("incomplete personal model listing facts")
+    return tuple((facts.names or ())[index] for index in decision.selected)
+
+
+async def _read(facts: PersonalModelListingFacts, source: _Source) -> PersonalModelListingFacts:
+    decision: Final = personal_model_listing(facts)
+    match decision.kind:
+        case "Complete":
+            return facts
+        case "ReadOwner":
+            owner: Final = await read_personal_owner_facts(
+                source.token, source.prisma, source.cache, source.logging, purpose="Models"
+            )
+            return await _read(replace(facts, owner=owner.facts), replace(source, teams=owner.teams))
+        case "ReadNames":
+            return await _read(replace(facts, names=tuple(source.names)), source)
+        case "ReadModel":
+            if source.prisma is None or source.cache is None or source.logging is None or facts.names is None:
+                raise NativeAuthRuntimeError("personal listing scope read without dependencies")
+            teams: Final = tuple(selected_team(source.teams, index) for index in decision.teams)
+            model: Final = await read_loaded_team_model_facts(
+                facts.names[decision.index],
+                teams,
+                source.token,
+                source.router,
+                source.prisma,
+                source.cache,
+                source.logging,
+            )
+            return await _read(replace(facts, models=(*facts.models, model)), source)

litellm/proxy/auth/personal_owner_adapter.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, replace
-from typing import TYPE_CHECKING, Final
+from typing import TYPE_CHECKING, Final, Literal
 
 from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError
 from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner
@@ -29,6 +29,7 @@ class _Source:
     logging: ProxyLogging | None
     check_db_only: bool
     owner: LiteLLM_UserTable | None = None
+    purpose: Literal["Models", "Grants"] = "Grants"
 
 
 async def read_personal_owner_facts(
@@ -38,6 +39,7 @@ async def read_personal_owner_facts(
     proxy_logging_obj: ProxyLogging | None,
     *,
     check_db_only: bool = False,
+    purpose: Literal["Models", "Grants"] = "Grants",
 ) -> PersonalOwnerReading:
     facts: Final = PersonalOwnerFacts(
         valid_token.user_id is not None,
@@ -47,7 +49,7 @@ async def read_personal_owner_facts(
     )
     return await _read(
         PersonalOwnerReading(facts),
-        _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only),
+        _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only, purpose=purpose),
     )
 
 
@@ -121,7 +123,7 @@ async def _read(reading: PersonalOwnerReading, source: _Source) -> PersonalOwner
                 source.prisma,
                 source.cache,
                 source.logging,
-                purpose="Grants",
+                purpose=source.purpose,
                 check_db_only=source.check_db_only,
             )
             return PersonalOwnerReading(replace(reading.facts, teams=teams.facts), teams.teams)

litellm/rust_bridge/_native.pyi

@@ -45,6 +45,9 @@ def request_auth_personal_owner(facts: object) -> object: ...
 def request_auth_personal_model_request(facts: object) -> object: ...
 def request_auth_member_models(facts: object) -> object: ...
 def request_auth_resolved_models(facts: object) -> object: ...
+def request_auth_personal_model_listing(facts: object) -> object: ...
+def request_auth_agent_ceiling(facts: object) -> object: ...
+def request_auth_agent_group_load(facts: object) -> object: ...
 
 @final
 class NativeClickHouseSpendConfig:

tests/unit/_auth/test_agent_ceilings.py

@@ -0,0 +1,91 @@
+from collections.abc import Callable
+from itertools import chain
+from typing import Final
+
+import pytest
+
+from litellm._auth.agent_ceilings import AgentCeilingDecision as Decision
+from litellm._auth.agent_ceilings import AgentCeilingFacts as Facts
+from litellm._auth.agent_ceilings import AgentCeilingGroup as Group
+from litellm._auth.agent_ceilings import AgentGroupLoadDecision as LoadDecision
+from litellm._auth.agent_ceilings import AgentGroupLoadFacts as LoadFacts
+from litellm._auth.agent_ceilings import agent_ceiling, agent_group_load
+from litellm._auth.failure import NativeAuthFailure
+
+
+def reference(facts: Facts) -> Decision:
+    if facts.ids is None:
+        return Decision("ReadIds")
+    if not facts.ids:
+        return Decision("Complete")
+    if facts.groups is None:
+        return Decision("ReadGroups", authoritative=facts.check_db_only)
+    indices: Final = tuple(index for index, group in enumerate(facts.groups) if group.status == "Found")
+    return Decision(
+        "Complete",
+        present=True,
+        ids=tuple(name for name, group in zip(facts.ids, facts.groups, strict=True) if group.status != "Missing"),
+        models=frozenset(chain.from_iterable(facts.groups[index].models for index in indices)),
+        mcp_server_ids=frozenset(chain.from_iterable(facts.groups[index].mcp_server_ids for index in indices)),
+        agent_ids=frozenset(chain.from_iterable(facts.groups[index].agent_ids for index in indices)),
+        indices=indices,
+    )
+
+
+def load_reference(facts: LoadFacts) -> LoadDecision:
+    if not facts.database_present:
+        return LoadDecision("DatabaseUnavailable")
+    if facts.status is None:
+        return LoadDecision("Read")
+    if facts.status == "Success":
+        return LoadDecision("Return")
+    return LoadDecision("Raise" if facts.status == "NotFound" or facts.check_db_only else "HttpUnavailable")
+
+
+@pytest.mark.parametrize("policy", (reference, agent_ceiling))
+@pytest.mark.parametrize(
+    ("facts", "expected"),
+    (
+        (Facts(), Decision("ReadIds")),
+        (Facts(()), Decision("Complete")),
+        (Facts(("group",)), Decision("ReadGroups")),
+        (Facts(("group",), check_db_only=True), Decision("ReadGroups", authoritative=True)),
+        (Facts(("group",), (Group("Missing"),)), Decision("Complete", present=True)),
+        (Facts(("group",), (Group("Unavailable"),)), Decision("Complete", present=True, ids=("group",))),
+        (
+            Facts(
+                ("one", "two", "one"),
+                (Group("Found", ("model",), ("server",), ("agent",)), Group("Missing"), Group("Found", ("model",))),
+            ),
+            Decision(
+                "Complete",
+                present=True,
+                ids=("one", "one"),
+                models=frozenset(("model",)),
+                mcp_server_ids=frozenset(("server",)),
+                agent_ids=frozenset(("agent",)),
+                indices=(0, 2),
+            ),
+        ),
+    ),
+)
+def test_ceiling_union_and_missing_vs_unavailable(
+    policy: Callable[[Facts], Decision], facts: Facts, expected: Decision
+) -> None:
+    assert policy(facts) == expected
+
+
+@pytest.mark.parametrize("policy", (load_reference, agent_group_load))
+@pytest.mark.parametrize("database", (False, True))
+@pytest.mark.parametrize("strict", (False, True))
+@pytest.mark.parametrize("status", (None, "Success", "NotFound", "HttpError"))
+def test_group_read_error_policy(
+    policy: Callable[[LoadFacts], LoadDecision], database: bool, strict: bool, status: str | None
+) -> None:
+    facts: Final = LoadFacts(database, strict, status)  # pyright: ignore[reportArgumentType]  # parameter table covers the complete typed status domain
+    assert policy(facts) == load_reference(facts)
+
+
+def test_mismatched_group_answers_are_terminal() -> None:
+    with pytest.raises(NativeAuthFailure):
+        agent_ceiling(Facts(("group",), ()))

tests/unit/_auth/test_personal_model_listing.py

@@ -0,0 +1,66 @@
+from collections.abc import Callable
+from dataclasses import replace
+from typing import Final
+
+import pytest
+
+from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad
+from litellm._auth.personal_model_listing import PersonalModelListingDecision as Decision
+from litellm._auth.personal_model_listing import PersonalModelListingFacts as Facts
+from litellm._auth.personal_model_listing import personal_model_listing
+from litellm._auth.personal_models import LoadedTeamModelFacts
+from litellm._auth.personal_owner import PersonalOwnerFacts
+from tests.unit._auth.test_personal_models import ALLOW, INITIAL
+from tests.unit._auth.test_personal_models import reference as scope_reference
+from tests.unit._auth.test_personal_owner import reference as owner_reference
+
+
+def reference(facts: Facts) -> Decision:
+    owner: Final = owner_reference(facts.owner)
+    if owner.kind != "Complete":
+        return Decision("ReadOwner")
+    if facts.owner.teams is None or not facts.owner.teams.loads:
+        return Decision("Complete")
+    if facts.names is None:
+        return Decision("ReadNames")
+    if len(facts.models) < len(facts.names):
+        return Decision("ReadModel", index=len(facts.models), teams=owner.selected)
+    return Decision(
+        "Complete",
+        selected=tuple(
+            index
+            for index, model in enumerate(facts.models)
+            if any(tuple(scope_reference(scope).allowed for scope in model.scopes))
+        ),
+    )
+
+
+_BASE: Final = Facts(PersonalOwnerFacts(True, True, True, True))
+_OWNER: Final = replace(
+    _BASE.owner,
+    owner_loaded=True,
+    owner_present=True,
+    teams=OwnerTeamFacts("Models", ("team",), "ui", loaded=True, loads=(OwnerTeamLoad(True),)),
+)
+_LOADED: Final = replace(_BASE, owner=_OWNER, names=("first", "second", "first"))
+_ALLOWED: Final = LoadedTeamModelFacts((True,), (ALLOW,))
+_DENIED: Final = LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),))
+
+
+@pytest.mark.parametrize("policy", (reference, personal_model_listing))
+@pytest.mark.parametrize(
+    ("facts", "expected"),
+    (
+        (_BASE, Decision("ReadOwner")),
+        (replace(_BASE, owner=replace(_BASE.owner, database_present=False)), Decision("Complete")),
+        (replace(_BASE, owner=replace(_BASE.owner, owner_loaded=True, owner_failed=True)), Decision("Complete")),
+        (replace(_BASE, owner=_OWNER), Decision("ReadNames")),
+        (_LOADED, Decision("ReadModel", teams=(0,))),
+        (replace(_LOADED, models=(_DENIED,)), Decision("ReadModel", index=1, teams=(0,))),
+        (replace(_LOADED, models=(_ALLOWED, _DENIED, _ALLOWED)), Decision("Complete", selected=(0, 2))),
+    ),
+)
+def test_listing_preserves_order_duplicates_and_checks_after_denial(
+    policy: Callable[[Facts], Decision], facts: Facts, expected: Decision
+) -> None:
+    assert policy(facts) == expected