spike: prototype the team_id column + (team_id, model_name) router identity

Exploratory spike used to stress-test the Team BYOK team_id plan against a live
proxy with mock upstreams. Not a shippable change: 5 tests in TestDeleteTeamModels
fail because they assert the prefix-scan query shape, and the v2 path regresses
retries for every team model.

- migration + schema: nullable indexed team_id on LiteLLM_ProxyModelTable (DDL only,
  no backfill - a backfill UPDATE is banned by check_migrations_no_data_rewrites)
- ownership queries read the column with a legacy prefix fallback
- rename guard reads the stored row instead of the caller payload
- get_model_group(id) scopes a team deployment's group to its owning team
- LITELLM_TEAM_MODEL_IDENTITY_V2=1: store the public name verbatim, keep team rows
  out of the global bare-name index and out of router.model_names

litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_proxy_model_team_id/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_ProxyModelTable" ADD COLUMN IF NOT EXISTS "team_id" TEXT;
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_ProxyModelTable_team_id_idx" ON "LiteLLM_ProxyModelTable"("team_id");

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -50,6 +50,7 @@ model LiteLLM_CredentialsTable {
 model LiteLLM_ProxyModelTable {
   model_id String @id @default(uuid())
   model_name String
+  team_id String?
   litellm_params Json
   model_info Json?
   blocked Boolean @default(false)
@@ -57,6 +58,8 @@ model LiteLLM_ProxyModelTable {
   created_by String
   updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
   updated_by String
+
+  @@index([team_id])
 }
 
 

litellm/models/model.py

@@ -16,6 +16,7 @@ from litellm.types.llms.base import LiteLLMPydanticObjectBase
 class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase):
     model_id: str
     model_name: str
+    team_id: str | None = None
     litellm_params: dict
     model_info: dict | None = None
     blocked: bool = False
@@ -48,7 +49,10 @@ class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase):
         return self.blocked
 
     @property
-    def team_id(self) -> str | None:
+    def owner_team_id(self) -> str | None:
+        """Team that owns this deployment: the column, falling back to the legacy JSON blob."""
+        if self.team_id is not None:
+            return self.team_id
         if self.model_info:
             return self.model_info.get("team_id")
         return None

litellm/proxy/_types.py

@@ -4821,6 +4821,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
 
 class PrismaCompatibleUpdateDBModel(TypedDict, total=False):
     model_name: str
+    team_id: str | None
     litellm_params: str
     model_info: str
     blocked: bool

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -33,6 +33,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
     SEARCH_CONTEXT_SIZES,
     ptu_config_error,
 )
+from litellm.team_model_identity import team_model_identity_v2_enabled
 from litellm.proxy._types import (
     BlockModelRequest,
     CommonProxyErrors,
@@ -603,6 +604,7 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
 
     prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel(
         model_name=merged_model_name,
+        team_id=merged_model_info.get("team_id"),
         litellm_params=json.dumps(merged_litellm_params),
         model_info=json.dumps(merged_model_info),
     )
@@ -962,6 +964,7 @@ async def _add_model_to_db(
     _data: Final[dict] = {
         "model_id": model_params.model_info.id,
         "model_name": model_params.model_name,
+        "team_id": model_params.model_info.team_id,
         "litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True),
         "model_info": model_params.model_info.model_dump_json(exclude_none=True),
         "created_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
@@ -1004,8 +1007,9 @@ async def _add_team_model_to_db(
 
     # Generate and assign unique internal model_name LAST
     # (after team_public_model_name is safely stored)
-    unique_model_name: Final = f"model_name_{_team_id}_{uuid.uuid4()}"
-    model_params.model_name = unique_model_name
+    if not team_model_identity_v2_enabled():
+        unique_model_name: Final = f"model_name_{_team_id}_{uuid.uuid4()}"
+        model_params.model_name = unique_model_name
 
     ## CREATE MODEL IN DB ##
     model_response: Final = await _add_model_to_db(
@@ -1069,9 +1073,13 @@ async def _update_team_model_in_db(
         ),
     )
 
-    patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None
+    # Ownership comes from the stored row, never from the caller's payload: a PATCH
+    # that sends only `model_name` used to take the standard-update path and rewrite
+    # a team deployment's routing key while `team_id` stayed behind, orphaning the row.
+    db_team_id_for_guard: Final = db_model.model_info.team_id if db_model.model_info else None
+    patch_team_id: Final = (patch_data.model_info.team_id if patch_data.model_info else None) or db_team_id_for_guard
 
-    # No team_id in patch, proceed with standard update
+    # No team ownership anywhere, proceed with standard update
     if patch_team_id is None:
         return update_db_model(db_model=db_model, updated_patch=patch_data)
 
@@ -1165,8 +1173,11 @@ async def _setup_new_team_model_assignment(
     user_api_key_dict: UserAPIKeyAuth,
 ) -> None:
     """Set up a new team model with unique name and team membership."""
-    unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}"
-    patch_data.model_name = unique_model_name
+    if not team_model_identity_v2_enabled():
+        unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}"
+        patch_data.model_name = unique_model_name
+    else:
+        patch_data.model_name = public_model_name
 
     await team_model_add(
         data=TeamModelAddRequest(
@@ -1196,17 +1207,25 @@ async def _get_team_deployments(
     """
     prefix: Final = f"model_name_{team_id}_"
     table = table or _proxy_model_table(prisma_client)
+    # The team_id column is the ownership key. The legacy name prefix is still read
+    # so rows written before the column existed (or by a pod that predates it) are
+    # not orphaned; the JSON check then confirms ownership for those.
     response: Final = await table.find_many(
         where={
-            "model_name": {"startswith": prefix},
+            "OR": [
+                {"team_id": team_id},
+                {"model_name": {"startswith": prefix}},
+            ]
         }
     )
     if not response:
         return []
 
-    # Confirm team_id in model_info (defensive check)
     result: Final = []
     for row in response:
+        if getattr(row, "team_id", None) == team_id:
+            result.append(row)
+            continue
         model_info = model_info_as_mapping(row.model_info)
         if model_info is not None and model_info.get("team_id") == team_id:
             result.append(row)
@@ -1383,7 +1402,9 @@ async def _update_existing_team_model_assignment(
     if old_public_name and public_model_name != old_public_name:
         # Clear user-supplied public name from patch before any early return so the
         # caller does not overwrite the internal UUID-based model_name in the DB.
-        patch_data.model_name = None
+        # Under the v2 identity the public name IS the stored name, so a rename must
+        # write it through instead of suppressing it.
+        patch_data.model_name = public_model_name if team_model_identity_v2_enabled() else None
         if prisma_client is None:
             verbose_proxy_logger.warning(
                 "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries"
@@ -1432,8 +1453,9 @@ async def _update_existing_team_model_assignment(
     # No team_model_add/delete calls required; public name is already registered
 
     # Always clear patch_data.model_name to prevent caller from overwriting
-    # the internal UUID-based model_name in the DB with the user-supplied public name
-    patch_data.model_name = None
+    # the internal UUID-based model_name in the DB with the user-supplied public name.
+    # Under v2 there is no internal name: the public name is the stored name.
+    patch_data.model_name = public_model_name if team_model_identity_v2_enabled() else None
 
 
 class ModelManagementAuthChecks:

litellm/proxy/schema.prisma

@@ -50,6 +50,7 @@ model LiteLLM_CredentialsTable {
 model LiteLLM_ProxyModelTable {
   model_id String @id @default(uuid())
   model_name String
+  team_id String?
   litellm_params Json
   model_info Json?
   blocked Boolean @default(false)
@@ -57,6 +58,8 @@ model LiteLLM_ProxyModelTable {
   created_by String
   updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
   updated_by String
+
+  @@index([team_id])
 }
 
 

litellm/router.py

@@ -117,6 +117,7 @@ from litellm.router_utils.batch_utils import (
     replace_model_in_jsonl,
     should_replace_model_in_jsonl,
 )
+from litellm.team_model_identity import team_model_identity_v2_enabled
 from litellm.router_utils.client_initalization_utils import InitalizeCachedClient
 from litellm.router_utils.clientside_credential_handler import (
     get_dynamic_litellm_params,
@@ -8798,7 +8799,11 @@ class Router:
                 )
 
         verbose_router_logger.debug("\nInitialized Model List %s", self.get_model_names())
-        self.model_names = {m["model_name"] for m in model_list}
+        self.model_names = {
+            m["model_name"]
+            for m in model_list
+            if not (team_model_identity_v2_enabled() and (m.get("model_info") or {}).get("team_id") is not None)
+        }
 
         # Note: model_name_to_deployment_indices is already built incrementally
         # by _create_deployment -> _add_model_to_list_and_index_map
@@ -9020,7 +9025,8 @@ class Router:
 
         # add to model names
         self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id)
-        self.model_names.add(deployment.model_name)
+        if not (team_model_identity_v2_enabled() and deployment.model_info.team_id is not None):
+            self.model_names.add(deployment.model_name)
         self._sync_deployment_budget_config(deployment=deployment)
         return deployment
 
@@ -9101,9 +9107,18 @@ class Router:
         - idx: int - the index in model_list
         """
         team_id: Final = (model.get("model_info") or {}).get("team_id")
-        team_public_model_name: Final = (model.get("model_info") or {}).get("team_public_model_name")
-        if team_id and team_public_model_name:
-            key: Final = (team_id, team_public_model_name)
+        if not team_id:
+            return
+        public_names: Final = {
+            name
+            for name in (
+                (model.get("model_info") or {}).get("team_public_model_name"),
+                model.get("model_name") if team_model_identity_v2_enabled() else None,
+            )
+            if name
+        }
+        for team_public_model_name in public_names:
+            key = (team_id, team_public_model_name)
             self.team_public_model_names = self.team_public_model_names | frozenset({team_public_model_name})
             if key not in self.team_model_to_deployment_indices:
                 self.team_model_to_deployment_indices[key] = []
@@ -9129,9 +9144,13 @@ class Router:
         elif model.get("model_info", {}).get("id") is not None:
             self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx
 
-        # Update model_name index for O(1) lookup
+        # Update model_name index for O(1) lookup.
+        # Under the v2 identity a team-owned deployment lives only under
+        # (team_id, model_name), so a lookup that forgets the team finds nothing
+        # instead of crossing into another team.
         model_name: Final = model.get("model_name")
-        if model_name:
+        owns_team: Final = (model.get("model_info") or {}).get("team_id") is not None
+        if model_name and not (team_model_identity_v2_enabled() and owns_team):
             if model_name not in self.model_name_to_deployment_indices:
                 self.model_name_to_deployment_indices[model_name] = []
             self.model_name_to_deployment_indices[model_name].append(idx)
@@ -9800,7 +9819,11 @@ class Router:
 
     def get_model_group(self, id: str) -> list | None:
         """
-        Return list of all models in the same model group as that model id
+        Return list of all models in the same model group as that model id.
+
+        A team-owned deployment's group is scoped to its owning team: the group is
+        ``(team_id, model_name)``, not the bare name, so a team model with sibling
+        replicas is not mistaken for a single-deployment group.
         """
 
         model_info: Final = self.get_model_info(id=id)
@@ -9808,6 +9831,10 @@ class Router:
             return None
 
         model_name: Final = model_info["model_name"]
+        owner_team_id: Final = (model_info.get("model_info") or {}).get("team_id")
+        if owner_team_id is not None:
+            public_name: Final = (model_info.get("model_info") or {}).get("team_public_model_name") or model_name
+            return self.get_model_list(model_name=public_name, team_id=owner_team_id)
         return self.get_model_list(model_name=model_name)
 
     def get_deployment_model_info(self, model_id: str, model_name: str) -> ModelInfo | None:
@@ -10541,6 +10568,8 @@ class Router:
             # Fallback: check by internal model_name for non-team deployments
             # or deployments that haven't been migrated to team_public_model_name yet
             model_team_id: Final = (model.get("model_info") or {}).get("team_id")
+            if team_model_identity_v2_enabled():
+                return model_team_id is None or model_team_id == team_id
             if (
                 team_id is None  # requester has no team constraint
                 or model_team_id is None  # global deployment - accessible to all teams
@@ -10640,6 +10669,11 @@ class Router:
                 team_model_name = self._get_team_specific_model(deployment=deployment, team_id=team_id)
                 if team_model_name:
                     model_names.append(team_model_name)
+                elif team_model_identity_v2_enabled() and team_id is None:
+                    # Admin listing: a team deployment still has a real public name.
+                    public = (model_info or {}).get("team_public_model_name") or deployment.get("model_name")
+                    if public:
+                        model_names.append(public)
             else:
                 model_names.append(deployment.get("model_name", ""))
 

litellm/team_model_identity.py

@@ -0,0 +1,16 @@
+"""Feature flag for the (team_id, model_name) deployment identity.
+
+v1 stores a team deployment under a synthetic ``model_name_{team_id}_{uuid}`` and
+translates back at every read site. v2 stores the public name verbatim and makes
+``(team_id or None, model_name)`` the router's key.
+"""
+
+import os
+from typing import Final
+
+
+def team_model_identity_v2_enabled() -> bool:
+    return os.getenv("LITELLM_TEAM_MODEL_IDENTITY_V2", "").lower() in ("1", "true", "yes")
+
+
+TEAM_MODEL_NAME_PREFIX: Final = "model_name_"

schema.prisma

@@ -50,6 +50,7 @@ model LiteLLM_CredentialsTable {
 model LiteLLM_ProxyModelTable {
   model_id String @id @default(uuid())
   model_name String
+  team_id String?
   litellm_params Json
   model_info Json?
   blocked Boolean @default(false)
@@ -57,6 +58,8 @@ model LiteLLM_ProxyModelTable {
   created_by String
   updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
   updated_by String
+
+  @@index([team_id])
 }