limited role api keys (#3115)

* in progress PoC

* working limited user, needs routes to be marked next

* make selected endpoint available to limited user role

* xfail on test_slack_prune

* add comment to sync function

---------

Co-authored-by: Richard Kuo <rkuo@rkuo.com>

Author: rkuo-danswer

backend/danswer/auth/schemas.py

@@ -15,6 +15,7 @@ class UserRole(str, Enum):
         for all groups they are a member of
     """
 
+    LIMITED = "limited"
     BASIC = "basic"
     ADMIN = "admin"
     CURATOR = "curator"

backend/danswer/auth/users.py

@@ -662,12 +662,26 @@ async def current_user_with_expired_token(
     return await double_check_user(user, include_expired=True)
 
 
-async def current_user(
+async def current_limited_user(
     user: User | None = Depends(optional_user),
 ) -> User | None:
     return await double_check_user(user)
 
 
+async def current_user(
+    user: User | None = Depends(optional_user),
+) -> User | None:
+    user = await double_check_user(user)
+    if not user:
+        return None
+
+    if user.role == UserRole.LIMITED:
+        raise BasicAuthenticationError(
+            detail="Access denied. User role is LIMITED. BASIC or higher permissions are required.",
+        )
+    return user
+
+
 async def current_curator_or_admin_user(
     user: User | None = Depends(current_user),
 ) -> User | None:

backend/danswer/server/auth_check.py

@@ -6,6 +6,7 @@
 
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_curator_or_admin_user
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.auth.users import current_user_with_expired_token
 from danswer.configs.app_configs import APP_API_PREFIX
@@ -102,7 +103,8 @@ def check_router_auth(
             for dependency in route_dependant_obj.dependencies:
                 depends_fn = dependency.cache_key[0]
                 if (
-                    depends_fn == current_user
+                    depends_fn == current_limited_user
+                    or depends_fn == current_user
                     or depends_fn == current_admin_user
                     or depends_fn == current_curator_or_admin_user
                     or depends_fn == api_key_dep
@@ -118,5 +120,5 @@ def check_router_auth(
             # print(f"(\"{route.path}\", {set(route.methods)}),")
 
             raise RuntimeError(
-                f"Did not find current_user or current_admin_user dependency in route - {route}"
+                f"Did not find user dependency in private route - {route}"
             )

backend/danswer/server/features/persona/api.py

@@ -11,6 +11,7 @@
 
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_curator_or_admin_user
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.configs.constants import FileOrigin
 from danswer.configs.constants import NotificationType
@@ -272,7 +273,7 @@ def list_personas(
 @basic_router.get("/{persona_id}")
 def get_persona(
     persona_id: int,
-    user: User | None = Depends(current_user),
+    user: User | None = Depends(current_limited_user),
     db_session: Session = Depends(get_session),
 ) -> PersonaSnapshot:
     return PersonaSnapshot.from_model(

backend/danswer/server/query_and_chat/chat_backend.py

@@ -18,6 +18,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.chat.chat_utils import create_chat_chain
 from danswer.chat.chat_utils import extract_headers
@@ -309,7 +310,7 @@ def is_connected_sync() -> bool:
 def handle_new_chat_message(
     chat_message_req: CreateChatMessageRequest,
     request: Request,
-    user: User | None = Depends(current_user),
+    user: User | None = Depends(current_limited_user),
     _: None = Depends(check_token_rate_limits),
     is_connected_func: Callable[[], bool] = Depends(is_connected),
 ) -> StreamingResponse:
@@ -391,7 +392,7 @@ def set_message_as_latest(
 @router.post("/create-chat-message-feedback")
 def create_chat_feedback(
     feedback: ChatFeedbackRequest,
-    user: User | None = Depends(current_user),
+    user: User | None = Depends(current_limited_user),
     db_session: Session = Depends(get_session),
 ) -> None:
     user_id = user.id if user else None

backend/danswer/server/query_and_chat/query_backend.py

@@ -9,6 +9,7 @@
 from sqlalchemy.orm import Session
 
 from danswer.auth.users import current_curator_or_admin_user
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.configs.constants import DocumentSource
 from danswer.configs.constants import MessageType
@@ -262,7 +263,7 @@ def stream_query_validation(
 @basic_router.post("/stream-answer-with-quote")
 def get_answer_with_quote(
     query_request: DirectQARequest,
-    user: User = Depends(current_user),
+    user: User = Depends(current_limited_user),
     _: None = Depends(check_token_rate_limits),
 ) -> StreamingResponse:
     query = query_request.messages[0].message

backend/tests/integration/common_utils/managers/cc_pair.py

@@ -326,6 +326,10 @@ def sync(
         cc_pair: DATestCCPair,
         user_performing_action: DATestUser | None = None,
     ) -> None:
+        """This function triggers a permission sync.
+        Naming / intent of this function probably could use improvement, but currently it's letting
+        409 Conflict pass through since if it's running that's what we were trying to do anyway.
+        """
         result = requests.post(
             url=f"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/sync-permissions",
             headers=user_performing_action.headers

backend/tests/integration/connector_job_tests/slack/test_prune.py

@@ -3,6 +3,8 @@
 from datetime import timezone
 from typing import Any
 
+import pytest
+
 from danswer.connectors.models import InputType
 from danswer.db.enums import AccessType
 from danswer.server.documents.models import DocumentSource
@@ -22,7 +24,7 @@
 from tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager
 
 
-# @pytest.mark.xfail(reason="flaky - see DAN-835 for example", strict=False)
+@pytest.mark.xfail(reason="flaky - see DAN-986 for details", strict=False)
 def test_slack_prune(
     reset: None,
     vespa_client: vespa_fixture,

backend/tests/integration/tests/api_key/test_api_key.py

@@ -0,0 +1,42 @@
+import requests
+
+from danswer.auth.schemas import UserRole
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.managers.api_key import APIKeyManager
+from tests.integration.common_utils.managers.user import UserManager
+from tests.integration.common_utils.test_models import DATestAPIKey
+from tests.integration.common_utils.test_models import DATestUser
+
+
+def test_limited(reset: None) -> None:
+    """Verify that with a limited role key, limited endpoints are accessible and
+    others are not."""
+
+    # Creating an admin user (first user created is automatically an admin)
+    admin_user: DATestUser = UserManager.create(name="admin_user")
+
+    api_key: DATestAPIKey = APIKeyManager.create(
+        api_key_role=UserRole.LIMITED,
+        user_performing_action=admin_user,
+    )
+
+    # test limited endpoint
+    response = requests.get(
+        f"{API_SERVER_URL}/persona/0",
+        headers=api_key.headers,
+    )
+    assert response.status_code == 200
+
+    # test basic endpoints
+    response = requests.get(
+        f"{API_SERVER_URL}/input_prompt",
+        headers=api_key.headers,
+    )
+    assert response.status_code == 403
+
+    # test admin endpoints
+    response = requests.get(
+        f"{API_SERVER_URL}/admin/api-key",
+        headers=api_key.headers,
+    )
+    assert response.status_code == 403

web/src/app/admin/api-key/DanswerApiKeyForm.tsx

@@ -2,14 +2,15 @@ import { Form, Formik } from "formik";
 import { PopupSpec } from "@/components/admin/connectors/Popup";
 import {
   BooleanFormField,
+  SelectorFormField,
   TextFormField,
 } from "@/components/admin/connectors/Field";
 import { createApiKey, updateApiKey } from "./lib";
 import { Modal } from "@/components/Modal";
 import { Button } from "@/components/ui/button";
 import { Separator } from "@/components/ui/separator";
 import Text from "@/components/ui/text";
-import { UserRole } from "@/lib/types";
+import { USER_ROLE_LABELS, UserRole } from "@/lib/types";
 import { APIKey } from "./types";
 
 interface DanswerApiKeyFormProps {
@@ -39,20 +40,15 @@ export const DanswerApiKeyForm = ({
         <Formik
           initialValues={{
             name: apiKey?.api_key_name || "",
-            is_admin: apiKey?.api_key_role === "admin",
+            role: apiKey?.api_key_role || UserRole.BASIC.toString(),
           }}
           onSubmit={async (values, formikHelpers) => {
             formikHelpers.setSubmitting(true);
 
-            // Map the boolean to a UserRole string
-            const role: UserRole = values.is_admin
-              ? UserRole.ADMIN
-              : UserRole.BASIC;
-
             // Prepare the payload with the UserRole
             const payload = {
               ...values,
-              role, // Assign the role directly as a UserRole type
+              role: values.role as UserRole, // Assign the role directly as a UserRole type
             };
 
             let response;
@@ -98,13 +94,28 @@ export const DanswerApiKeyForm = ({
                 autoCompleteDisabled={true}
               />
 
-              <BooleanFormField
-                small
-                removeIndent
-                alignTop
-                name="is_admin"
-                label="Is Admin?"
-                subtext="If set, this API key will have access to admin level server API's."
+              <SelectorFormField
+                // defaultValue is managed by Formik
+                label="Role:"
+                subtext="Select the role for this API key.
+                         Limited has access to simple public API's.
+                         Basic has access to regular user API's.
+                         Admin has access to admin level APIs."
+                name="role"
+                options={[
+                  {
+                    name: USER_ROLE_LABELS[UserRole.LIMITED],
+                    value: UserRole.LIMITED.toString(),
+                  },
+                  {
+                    name: USER_ROLE_LABELS[UserRole.BASIC],
+                    value: UserRole.BASIC.toString(),
+                  },
+                  {
+                    name: USER_ROLE_LABELS[UserRole.ADMIN],
+                    value: UserRole.ADMIN.toString(),
+                  },
+                ]}
               />
 
               <Button