fix: attach team to org table (#17832)

* fix: attach team to org table

* add test

litellm/proxy/management_endpoints/team_endpoints.py

@@ -1052,7 +1052,7 @@ async def fetch_and_validate_organization(
 
     organization_row = await prisma_client.db.litellm_organizationtable.find_unique(
         where={"organization_id": organization_id},
-        include={"litellm_budget_table": True, "members": True},
+        include={"litellm_budget_table": True, "members": True, "teams": True},
     )
 
     if organization_row is None:

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -3726,4 +3726,139 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit():
         )
 
         # Verify team was updated
-        assert result["team_id"] == "org-team-update-bypass-123"
+        assert result["team_id"] == "org-team-update-bypass-123"
+
+
+@pytest.mark.asyncio
+async def test_update_team_guardrails_with_org_id():
+    """
+    Test that updating team guardrails works when team has an organization_id.
+    The fix ensures 'teams' field is included when fetching organization data.
+    """
+    from fastapi import Request
+
+    from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_TeamTable
+    from litellm.proxy.management_endpoints.team_endpoints import update_team
+
+    # Create user (org admin)
+    org_admin_user = UserAPIKeyAuth(
+        user_role=LitellmUserRoles.INTERNAL_USER,
+        user_id="org-admin-guardrails-test",
+        models=[],
+    )
+
+    # Update request to add guardrails to team
+    update_request = UpdateTeamRequest(
+        team_id="team-guardrails-123",
+        guardrails=["aporia-pre-call", "aporia-post-call"],
+        organization_id="test-org-guardrails",  # Changing org triggers fetch_and_validate_organization
+    )
+
+    dummy_request = MagicMock(spec=Request)
+
+    # Mock organization with all required fields including teams (the fix)
+    from datetime import datetime
+    mock_org = MagicMock(spec=LiteLLM_OrganizationTable)
+    mock_org.organization_id = "test-org-guardrails"
+    mock_org.models = ["gpt-4", "gpt-3.5-turbo"]
+    mock_org.budget_id = "budget-123"
+    mock_org.created_by = "admin"
+    mock_org.updated_by = "admin"
+    mock_org.created_at = datetime(2024, 1, 1)
+    mock_org.updated_at = datetime(2024, 1, 1)
+    mock_org.litellm_budget_table = None
+    mock_org.members = []
+    mock_org.teams = []  # Must be a list, not None
+    mock_org.model_dump.return_value = {
+        "organization_id": "test-org-guardrails",
+        "models": ["gpt-4", "gpt-3.5-turbo"],
+        "budget_id": "budget-123",
+        "created_by": "admin",
+        "updated_by": "admin",
+        "created_at": datetime(2024, 1, 1),
+        "updated_at": datetime(2024, 1, 1),
+        "litellm_budget_table": None,
+        "members": [],
+        "teams": [],
+    }
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
+        "litellm.proxy.proxy_server.user_api_key_cache"
+    ) as mock_cache, patch(
+        "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"
+    ), patch(
+        "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()
+    ):
+        # Mock existing team - must have compatible models with organization
+        mock_existing_team = MagicMock()
+        mock_existing_team.team_id = "team-guardrails-123"
+        mock_existing_team.organization_id = None
+        mock_existing_team.metadata = {}
+        mock_existing_team.model_id = None
+        mock_existing_team.models = ["gpt-4"]  # Subset of org models to pass validation
+        mock_existing_team.max_budget = None
+        mock_existing_team.tpm_limit = None
+        mock_existing_team.rpm_limit = None
+        mock_existing_team.model_dump.return_value = {
+            "team_id": "team-guardrails-123",
+            "organization_id": None,
+            "metadata": {},
+            "models": ["gpt-4"],
+            "max_budget": None,
+            "tpm_limit": None,
+            "rpm_limit": None,
+        }
+        mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
+            return_value=mock_existing_team
+        )
+        mock_cache.async_set_cache = AsyncMock()
+
+        # Mock organization fetch - this is where the bug occurred
+        # The fix ensures 'teams: True' is in the include clause
+        mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock(
+            return_value=mock_org
+        )
+
+        # Mock team update
+        mock_updated_team = MagicMock(spec=LiteLLM_TeamTable)
+        mock_updated_team.team_id = "team-guardrails-123"
+        mock_updated_team.organization_id = "test-org-guardrails"
+        mock_updated_team.metadata = {"guardrails": ["aporia-pre-call", "aporia-post-call"]}
+        mock_updated_team.litellm_model_table = None
+        mock_updated_team.model_dump.return_value = {
+            "team_id": "team-guardrails-123",
+            "organization_id": "test-org-guardrails",
+            "metadata": {"guardrails": ["aporia-pre-call", "aporia-post-call"]},
+        }
+        mock_prisma.db.litellm_teamtable.update = AsyncMock(
+            return_value=mock_updated_team
+        )
+        mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
+
+        # Mock llm_router
+        mock_router = MagicMock()
+        with patch("litellm.proxy.proxy_server.llm_router", mock_router):
+            # This should succeed without Pydantic validation error
+            result = await update_team(
+                data=update_request,
+                http_request=dummy_request,
+                user_api_key_dict=org_admin_user,
+            )
+
+            # Verify the team was updated successfully with guardrails
+            assert result is not None
+            assert result["data"].organization_id == "test-org-guardrails"
+            assert result["data"].metadata["guardrails"] == ["aporia-pre-call", "aporia-post-call"]
+
+            # Verify that organization fetch was called with proper include clause
+            # The function is called twice: once by fetch_and_validate_organization (with include)
+            # and once by get_org_object (without include). We verify the first call has 'teams'.
+            assert mock_prisma.db.litellm_organizationtable.find_unique.call_count >= 1
+            
+            # Get the first call (from fetch_and_validate_organization)
+            first_call_kwargs = mock_prisma.db.litellm_organizationtable.find_unique.call_args_list[0].kwargs
+            
+            # Verify that 'teams' is included in the fetch
+            assert "include" in first_call_kwargs
+            assert "teams" in first_call_kwargs["include"]
+            assert first_call_kwargs["include"]["teams"] is True