Fix Agent bot 2 permission bugs (#818) (#16909)

api/apps/restful_apis/agent_api.py

@@ -35,7 +35,7 @@ from api.utils.file_response import (
 import jwt
 from quart import Response, jsonify, request, make_response
 
-from api.apps import current_user, login_required
+from api.apps import AUTH_JWT, AUTH_API, AUTH_BETA, current_user, login_required
 from api.apps.services.canvas_replica_service import CanvasReplicaService
 from api.db import CanvasCategory
 from api.db.db_models import Task
@@ -2514,7 +2514,7 @@ async def preview_attachment(tenant_id=None, attachment_id=None):
 
 
 @manager.route("/agents/attachments/<attachment_id>/download", methods=["GET"])  # noqa: F821
-@login_required
+@login_required(auth_types=[AUTH_JWT, AUTH_API, AUTH_BETA])
 @add_tenant_id_to_kwargs
 async def download_attachment(tenant_id=None, attachment_id=None):
     """Stream an agent-generated attachment as a download."""

api/apps/restful_apis/bot_api.py

@@ -14,7 +14,6 @@
 #  limitations under the License.
 #
 import copy
-import hashlib
 import json
 import re
 
@@ -24,7 +23,6 @@ from quart import Response, request
 
 from agent.canvas import Canvas
 from api.apps import AUTH_BETA, login_required
-from api.db.db_models import APIToken
 from api.db.services.api_service import API4ConversationService
 from api.db.services.canvas_service import UserCanvasService
 from api.db.services.canvas_service import completion as agent_completion
@@ -53,13 +51,6 @@ from api.utils.reference_metadata_utils import (
 logger = logging.getLogger(__name__)
 
 
-def _get_sdk_authorization_token():
-    auth_header = request.headers.get("Authorization", "")
-    if not auth_header.startswith("Bearer "):
-        return ""
-    return auth_header[len("Bearer ") :].strip()
-
-
 @manager.route("/chatbots/<dialog_id>/completions", methods=["POST"])  # noqa: F821
 @login_required(auth_types=AUTH_BETA)
 @add_tenant_id_to_kwargs
@@ -275,47 +266,20 @@ async def begin_inputs(agent_id, tenant_id=None):
 
 
 @manager.route("/agentbots/<shared_id>/logs/<message_id>", methods=["GET"])  # noqa: F821
-async def agent_bot_logs(shared_id, message_id):
-    # Beta-token sibling of /agents/<agent_id>/logs/<message_id>.
-    # Used by the shared/embedded chat page's "Thinking" button (fixes #14985).
-    # The <shared_id> path segment is just the value the client passed in the
-    # URL (it equals the beta token in the share flow); authentication comes
-    # from the Authorization header and the real agent_id is read from the
-    # looked-up APIToken so we never trust client-supplied identifiers.
+@login_required(auth_types=AUTH_BETA)
+@add_tenant_id_to_kwargs
+async def agent_bot_logs(shared_id, message_id, tenant_id=None):
+    if not await thread_pool_exec(UserCanvasService.accessible, shared_id, tenant_id):
+        logger.warning(
+            "agent bot logs access denied tenant_id=%s agent_id=%s",
+            tenant_id,
+            shared_id,
+        )
+        return get_error_data_result(f"Can't find agent by ID: {shared_id}")
     from rag.utils.redis_conn import REDIS_CONN
 
-    token = _get_sdk_authorization_token()
-    if not token:
-        logger.warning(
-            "agent_bot_logs: missing Authorization header (shared_id=%s message_id=%s)",
-            shared_id,
-            message_id,
-        )
-        return get_error_data_result(message="Authorization is not valid!")
-    # Non-reversible fingerprint of the share token: lets operators correlate
-    # auth-failure log lines for the same token without leaking a guessable
-    # substring of the secret itself.
-    token_fp = hashlib.sha256(token.encode("utf-8")).hexdigest()[:16]
-    objs = await thread_pool_exec(APIToken.query, beta=token)
-    if not objs:
-        logger.warning(
-            "agent_bot_logs: invalid beta token (fingerprint=%s shared_id=%s)",
-            token_fp,
-            shared_id,
-        )
-        return get_error_data_result(message='Authentication error: API key is invalid!"')
-
-    agent_id = objs[0].dialog_id
-    if not agent_id:
-        logger.warning(
-            "agent_bot_logs: APIToken has no dialog_id (tenant_id=%s fingerprint=%s)",
-            objs[0].tenant_id,
-            token_fp,
-        )
-        return get_error_data_result(message="API token is not bound to an agent.")
-
     try:
-        binary = await thread_pool_exec(REDIS_CONN.get, f"{agent_id}-{message_id}-logs")
+        binary = await thread_pool_exec(REDIS_CONN.get, f"{shared_id}-{message_id}-logs")
         if not binary:
             return get_json_result(data={})
         payload = binary.decode("utf-8") if isinstance(binary, bytes) else binary

test/unit_test/api/apps/restful_apis/test_attachment_download_missing_blob.py

@@ -71,8 +71,11 @@ def _load_agent_api(monkeypatch, *, storage_get):
     _stub(
         monkeypatch,
         "api.apps",
+        AUTH_JWT="jwt",
+        AUTH_API="api",
+        AUTH_BETA="beta",
         current_user=SimpleNamespace(id="tenant-1"),
-        login_required=lambda func: func,
+        login_required=lambda func=None, **_kwargs: (lambda f: f) if func is None else func,
     )
     _stub(monkeypatch, "api.apps.services.canvas_replica_service", CanvasReplicaService=SimpleNamespace())
     _stub(monkeypatch, "api.db", CanvasCategory=SimpleNamespace())

test/unit_test/api/apps/restful_apis/test_get_agent_session.py

@@ -44,7 +44,16 @@ def _load_agent_api(monkeypatch, get_by_id_result, delete_calls=None):
         delete_calls.append(session_id)
         return True
 
-    _stub(monkeypatch, "api.apps", current_user=SimpleNamespace(id="tenant-1"), login_required=lambda func: func)
+    _stub(
+        monkeypatch,
+        "api.apps",
+        AUTH_JWT="jwt",
+        AUTH_API="api",
+        AUTH_BETA="beta",
+        QuartAuthUnauthorized=Exception,
+        current_user=SimpleNamespace(id="tenant-1", is_superuser=True),
+        login_required=lambda func=None, **_kwargs: (lambda f: f) if func is None else func,
+    )
     _stub(monkeypatch, "api.apps.services.canvas_replica_service", CanvasReplicaService=SimpleNamespace())
     _stub(monkeypatch, "api.db", CanvasCategory=SimpleNamespace())
     _stub(monkeypatch, "api.db.db_models", Task=SimpleNamespace())

test/unit_test/api/apps/restful_apis/test_rerun_agent_authorization.py

@@ -48,7 +48,16 @@ def _load_agent_api_for_rerun(monkeypatch, *, documents_info, accessible):
 
     acc_fn = accessible if callable(accessible) else (lambda *_a, **_k: accessible)
 
-    _stub(monkeypatch, "api.apps", current_user=SimpleNamespace(id="user-owner"), login_required=lambda func: func)
+    _stub(
+        monkeypatch,
+        "api.apps",
+        AUTH_JWT="jwt",
+        AUTH_API="api",
+        AUTH_BETA="beta",
+        QuartAuthUnauthorized=Exception,
+        current_user=SimpleNamespace(id="user-owner", is_superuser=True),
+        login_required=lambda func=None, **_kwargs: (lambda f: f) if func is None else func,
+    )
     _stub(monkeypatch, "api.apps.services.canvas_replica_service", CanvasReplicaService=SimpleNamespace())
     _stub(monkeypatch, "api.db", CanvasCategory=SimpleNamespace())