Allow 'key_creator' users to fully manage team keys

Author: PhiRho

app/api/private_ai_keys.py

@@ -55,14 +55,14 @@ def _validate_permissions_and_get_ownership_info(
         owner_id = current_user.id
 
     # Fail fast without having to do DB lookups
-    personal_users : list[UserRole] = ["key_creator", "user"] # 'user' is the non-admin system user
-    if user_role in personal_users:
+    team_users : list[UserRole] = [UserRole.TEAM_ADMIN, UserRole.KEY_CREATOR]
+    if user_role == UserRole.DEFAULT:
         if owner_id != current_user.id or team_id is not None:
             raise HTTPException(
                 status_code=status.HTTP_403_FORBIDDEN,
                 detail="Not authorized to perform this action"
             )
-    elif user_role == "admin": # roles always refer to the team role, so admin is a team admin
+    elif user_role in team_users:
         if team_id is not None and team_id != current_user.team_id:
             raise HTTPException(
                 status_code=status.HTTP_403_FORBIDDEN,
@@ -455,7 +455,7 @@ async def list_private_ai_keys(
     else:
         # Check if user is a team admin
         if current_user.team_id is not None:
-            if current_user.role == "admin":
+            if current_user.role == UserRole.TEAM_ADMIN:
                 # Get all users in the team
                 team_users = db.query(DBUser).filter(DBUser.team_id == current_user.team_id).all()
                 team_user_ids = [user.id for user in team_users]
@@ -551,22 +551,25 @@ def _get_key_if_allowed(key_id: int, current_user: DBUser, user_role: UserRole,
         )
 
     # Check if user has permission to view the key
+    team_users : list[UserRole] = [UserRole.TEAM_ADMIN, UserRole.KEY_CREATOR]
     if current_user.is_admin:
         # System admin can view any key
         pass
-    elif user_role == "admin":
-        # Team admin can only view keys from their team
+    elif user_role in team_users:
+        # No cross team access
         if private_ai_key.team_id is not None:
             # For team-owned keys, check if it belongs to the admin's team
             if private_ai_key.team_id != current_user.team_id:
+                logger.warning(f"User {current_user.id} trying to delete across teams.")
                 raise HTTPException(
                     status_code=status.HTTP_404_NOT_FOUND,
                     detail="Private AI Key not found"
                 )
         else:
-            # For user-owned keys, check if the owner is in the admin's team
+            # For user-owned keys, check if the owner is in the viewer's team
             owner = db.query(DBUser).filter(DBUser.id == private_ai_key.owner_id).first()
             if not owner or owner.team_id != current_user.team_id:
+                logger.warning(f"Keys owned by different teams")
                 raise HTTPException(
                     status_code=status.HTTP_404_NOT_FOUND,
                     detail="Private AI Key not found"

app/core/resource_limits.py

@@ -3,17 +3,16 @@
 from app.db.models import DBTeam, DBUser, DBPrivateAIKey, DBTeamProduct, DBProduct
 from fastapi import HTTPException, status
 from typing import Optional
-from datetime import datetime, UTC
 import logging
 
 logger = logging.getLogger(__name__)
 
 # Default limits across all customers and products
-DEFAULT_USER_COUNT = 100
-DEFAULT_KEYS_PER_USER = 5
-DEFAULT_TOTAL_KEYS = 500
-DEFAULT_SERVICE_KEYS = 100
-DEFAULT_VECTOR_DB_COUNT = 100
+DEFAULT_USER_COUNT = 1
+DEFAULT_KEYS_PER_USER = 1
+DEFAULT_TOTAL_KEYS = 6
+DEFAULT_SERVICE_KEYS = 5
+DEFAULT_VECTOR_DB_COUNT = 5 # Setting to match service keys for drupal module trial
 DEFAULT_KEY_DURATION = 30
 DEFAULT_MAX_SPEND = 27.0
 DEFAULT_RPM_PER_KEY = 500

frontend/src/app/private-ai-keys/page.tsx

@@ -4,10 +4,10 @@ import { useState } from 'react';
 import { useMutation, useQueryClient, useQuery } from '@tanstack/react-query';
 import { Card, CardContent } from '@/components/ui/card';
 import { useToast } from '@/hooks/use-toast';
-import { get, post } from '@/utils/api';
+import { get, post, del } from '@/utils/api';
 import { PrivateAIKeysTable } from '@/components/private-ai-keys-table';
 import { CreateAIKeyDialog } from '@/components/create-ai-key-dialog';
-import { useAuth } from '@/hooks/use-auth';
+import { getUserRole, useAuth } from '@/hooks/use-auth';
 import { usePrivateAIKeysData } from '@/hooks/use-private-ai-keys-data';
 
 
@@ -32,6 +32,7 @@ export default function DashboardPage() {
   const queryClient = useQueryClient();
   const { user } = useAuth();
   const [isCreateDialogOpen, setIsCreateDialogOpen] = useState(false);
+  const mutatingRoles = ["admin", "key_creator", "system_admin"]
 
   // Fetch private AI keys using React Query
   const { data: privateAIKeys = [] } = useQuery<PrivateAIKey[]>({
@@ -76,7 +77,7 @@ export default function DashboardPage() {
   // Delete key mutation
   const deleteKeyMutation = useMutation({
     mutationFn: async (keyId: number) => {
-      const response = await post(`/private-ai-keys/${keyId}/delete`, {});
+      const response = await del(`/private-ai-keys/${keyId}`, {});
       return response.json();
     },
     onSuccess: () => {
@@ -107,8 +108,8 @@ export default function DashboardPage() {
       team_id?: number
     }) => {
       const endpoint = key_type === 'full' ? '/private-ai-keys' :
-                      key_type === 'llm' ? '/private-ai-keys/token' :
-                      '/private-ai-keys/vector-db';
+        key_type === 'llm' ? '/private-ai-keys/token' :
+          '/private-ai-keys/vector-db';
       const payload: { region_id: number; name: string; owner_id?: number; team_id?: number } = { region_id, name };
       if (owner_id) payload.owner_id = owner_id;
       if (team_id) payload.team_id = team_id;
@@ -176,7 +177,7 @@ export default function DashboardPage() {
         onDelete={(keyId) => deleteKeyMutation.mutate(keyId)}
         isLoading={createKeyMutation.isPending}
         showOwner={true}
-        allowModification={false}
+        allowModification={mutatingRoles.includes(getUserRole(user))}
         onUpdateBudget={(keyId, budgetDuration) => updateBudgetMutation.mutate({ keyId, budgetDuration })}
         isDeleting={deleteKeyMutation.isPending}
         isUpdatingBudget={updateBudgetMutation.isPending}

frontend/src/hooks/use-auth.ts

@@ -25,4 +25,11 @@ export const useAuth = create<AuthState>((set) => ({
 export const isTeamAdmin = (user: User | null): boolean => {
   if (!user) return false;
   return !user.is_admin && user.team_id !== null && user.role === 'admin';
-};
+};
+
+export const getUserRole = (user: User | null): string => {
+  if (!user) return "none";
+  if (user.is_admin) return "system_admin";
+  if (!user.role) return "user";
+  return user.role;
+}