feat(dify-cli): session level tool white list

api/controllers/cli_api/plugin/plugin.py

@@ -1,3 +1,4 @@
+from flask import abort
 from flask_restx import Resource
 
 from controllers.cli_api import cli_api_ns
@@ -15,6 +16,8 @@ from core.plugin.entities.request import (
     RequestInvokeTool,
     RequestRequestUploadFile,
 )
+from core.session.cli_api import CliContext
+from core.skill.entities import ToolInvocationRequest
 from core.tools.entities.tool_entities import ToolProviderType
 from libs.helper import length_prefixed_response
 from models import Account, Tenant
@@ -23,9 +26,9 @@ from models.model import EndUser
 
 @cli_api_ns.route("/invoke/llm")
 class CliInvokeLLMApi(Resource):
+    @cli_api_only
     @get_cli_user_tenant
     @setup_required
-    @cli_api_only
     @plugin_data(payload_type=RequestInvokeLLM)
     def post(self, user_model: Account | EndUser, tenant_model: Tenant, payload: RequestInvokeLLM):
         def generator():
@@ -37,17 +40,34 @@ class CliInvokeLLMApi(Resource):
 
 @cli_api_ns.route("/invoke/tool")
 class CliInvokeToolApi(Resource):
+    @cli_api_only
     @get_cli_user_tenant
     @setup_required
-    @cli_api_only
     @plugin_data(payload_type=RequestInvokeTool)
-    def post(self, user_model: Account | EndUser, tenant_model: Tenant, payload: RequestInvokeTool):
+    def post(
+        self,
+        user_model: Account | EndUser,
+        tenant_model: Tenant,
+        payload: RequestInvokeTool,
+        cli_context: CliContext,
+    ):
+        tool_type = ToolProviderType.value_of(payload.tool_type)
+
+        request = ToolInvocationRequest(
+            tool_type=tool_type,
+            provider=payload.provider,
+            tool_name=payload.tool,
+            credential_id=payload.credential_id,
+        )
+        if cli_context.tool_access and not cli_context.tool_access.is_allowed(request):
+            abort(403)
+
         def generator():
             return PluginToolBackwardsInvocation.convert_to_event_stream(
                 PluginToolBackwardsInvocation.invoke_tool(
                     tenant_id=tenant_model.id,
                     user_id=user_model.id,
-                    tool_type=ToolProviderType.value_of(payload.tool_type),
+                    tool_type=tool_type,
                     provider=payload.provider,
                     tool_name=payload.tool,
                     tool_parameters=payload.tool_parameters,
@@ -60,9 +80,9 @@ class CliInvokeToolApi(Resource):
 
 @cli_api_ns.route("/invoke/app")
 class CliInvokeAppApi(Resource):
+    @cli_api_only
     @get_cli_user_tenant
     @setup_required
-    @cli_api_only
     @plugin_data(payload_type=RequestInvokeApp)
     def post(self, user_model: Account | EndUser, tenant_model: Tenant, payload: RequestInvokeApp):
         response = PluginAppBackwardsInvocation.invoke_app(
@@ -81,9 +101,9 @@ class CliInvokeAppApi(Resource):
 
 @cli_api_ns.route("/upload/file/request")
 class CliUploadFileRequestApi(Resource):
+    @cli_api_only
     @get_cli_user_tenant
     @setup_required
-    @cli_api_only
     @plugin_data(payload_type=RequestRequestUploadFile)
     def post(self, user_model: Account | EndUser, tenant_model: Tenant, payload: RequestRequestUploadFile):
         # generate signed url
@@ -98,9 +118,9 @@ class CliUploadFileRequestApi(Resource):
 
 @cli_api_ns.route("/fetch/tools/list")
 class CliFetchToolsListApi(Resource):
+    @cli_api_only
     @get_cli_user_tenant
     @setup_required
-    @cli_api_only
     def post(self, user_model: Account | EndUser, tenant_model: Tenant):
         from sqlalchemy.orm import Session
 

api/controllers/cli_api/plugin/wraps.py

@@ -2,12 +2,12 @@ from collections.abc import Callable
 from functools import wraps
 from typing import ParamSpec, TypeVar
 
-from flask import current_app, request
+from flask import current_app, g, request
 from flask_login import user_logged_in
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
-from core.session.cli_api import CliApiSession, CliApiSessionManager
+from core.session.cli_api import CliApiSession, CliContext
 from extensions.ext_database import db
 from libs.login import current_user
 from models.account import Tenant
@@ -75,22 +75,13 @@ def get_user(tenant_id: str, user_id: str | None) -> EndUser:
 def get_cli_user_tenant(view_func: Callable[P, R]):
     @wraps(view_func)
     def decorated_view(*args: P.args, **kwargs: P.kwargs):
-        session_id = request.headers.get("X-Cli-Api-Session-Id")
+        session: CliApiSession | None = getattr(g, "cli_api_session", None)
+        if session is None:
+            raise ValueError("session not found")
 
-        if session_id:
-            session: CliApiSession | None = CliApiSessionManager().get(session_id)
-            if not session:
-                raise ValueError("session not found")
-            user_id = session.user_id
-            tenant_id = session.tenant_id
-
-        else:
-            payload = TenantUserPayload.model_validate(request.get_json(silent=True) or {})
-            user_id = payload.user_id
-            tenant_id = payload.tenant_id
-
-        if not tenant_id:
-            raise ValueError("tenant_id is required")
+        user_id = session.user_id
+        tenant_id = session.tenant_id
+        cli_context = CliContext.model_validate(session.context)
 
         if not user_id:
             user_id = DefaultEndUserSessionID.DEFAULT_SESSION_ID
@@ -110,11 +101,10 @@ def get_cli_user_tenant(view_func: Callable[P, R]):
             raise ValueError("tenant not found")
 
         kwargs["tenant_model"] = tenant_model
+        kwargs["user_model"] = get_user(tenant_id, user_id)
+        kwargs["cli_context"] = cli_context
 
-        user = get_user(tenant_id, user_id)
-        kwargs["user_model"] = user
-
-        current_app.login_manager._update_request_context_with_user(user)  # type: ignore
+        current_app.login_manager._update_request_context_with_user(kwargs["user_model"])  # type: ignore
         user_logged_in.send(current_app._get_current_object(), user=current_user)  # type: ignore
 
         return view_func(*args, **kwargs)