feat: session management for InnerAPI&VM

api/controllers/inner_api/plugin/wraps.py

@@ -7,6 +7,7 @@ from flask_login import user_logged_in
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from core.session.inner_api import InnerApiSession, InnerApiSessionManager
 from extensions.ext_database import db
 from libs.login import current_user
 from models.account import Tenant
@@ -74,10 +75,18 @@ def get_user(tenant_id: str, user_id: str | None) -> EndUser:
 def get_user_tenant(view_func: Callable[P, R]):
     @wraps(view_func)
     def decorated_view(*args: P.args, **kwargs: P.kwargs):
-        payload = TenantUserPayload.model_validate(request.get_json(silent=True) or {})
+        session_id = request.headers.get("X-Inner-Api-Session-Id")
 
-        user_id = payload.user_id
-        tenant_id = payload.tenant_id
+        if session_id:
+            session: InnerApiSession | None = InnerApiSessionManager().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")

api/controllers/inner_api/wraps.py

@@ -5,6 +5,8 @@ from hashlib import sha1
 from hmac import new as hmac_new
 from typing import ParamSpec, TypeVar
 
+from core.session.inner_api import InnerApiSessionManager
+
 P = ParamSpec("P")
 R = TypeVar("R")
 from flask import abort, request
@@ -85,14 +87,19 @@ def enterprise_inner_api_user_auth(view: Callable[P, R]):
 def plugin_inner_api_only(view: Callable[P, R]):
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
+        # if session id is provided, using session id to validate
+        session_id = request.headers.get("X-Inner-Api-Session-Id")
+        if session_id and InnerApiSessionManager().exists(session_id):
+            return view(*args, **kwargs)
+
         if not dify_config.PLUGIN_DAEMON_KEY:
             abort(404)
 
-        # get header 'X-Inner-Api-Key'
+        # if inner api key is provided, using inner api key to validate
         inner_api_key = request.headers.get("X-Inner-Api-Key")
-        if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY_FOR_PLUGIN:
-            abort(404)
+        if inner_api_key and inner_api_key == dify_config.INNER_API_KEY_FOR_PLUGIN:
+            return view(*args, **kwargs)
 
-        return view(*args, **kwargs)
+        abort(404)
 
     return decorated