From 2b97f6c8c224ed373709a158f9b0af3bb343f077 Mon Sep 17 00:00:00 2001
From: Tianle <40735546+Tianlel@users.noreply.github.com>
Date: Tue, 26 May 2026 00:45:03 -0500
Subject: [PATCH] chore: inject tenant id in extension handlers (#36656)

---
 api/controllers/console/extension.py          | 28 +++++++++----------
 .../controllers/console/test_extension.py     |  2 --
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/api/controllers/console/extension.py b/api/controllers/console/extension.py
index 029e2e7f0d..26a348da40 100644
--- a/api/controllers/console/extension.py
+++ b/api/controllers/console/extension.py
@@ -9,14 +9,14 @@ from pydantic import BaseModel, Field, TypeAdapter, field_validator
 from constants import HIDDEN_VALUE
 from fields.base import ResponseModel
 from libs.helper import to_timestamp
-from libs.login import current_account_with_tenant, login_required
+from libs.login import login_required
 from models.api_based_extension import APIBasedExtension
 from services.api_based_extension_service import APIBasedExtensionService
 from services.code_based_extension_service import CodeBasedExtensionService
 
 from ..common.schema import DEFAULT_REF_TEMPLATE_SWAGGER_2_0, register_schema_models
 from . import console_ns
-from .wraps import account_initialization_required, setup_required
+from .wraps import account_initialization_required, setup_required, with_current_tenant_id
 
 
 class CodeBasedExtensionQuery(BaseModel):
@@ -116,11 +116,11 @@ class APIBasedExtensionAPI(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    def get(self):
-        _, tenant_id = current_account_with_tenant()
+    @with_current_tenant_id
+    def get(self, current_tenant_id: str):
         return [
             _serialize_api_based_extension(extension)
-            for extension in APIBasedExtensionService.get_all_by_tenant_id(tenant_id)
+            for extension in APIBasedExtensionService.get_all_by_tenant_id(current_tenant_id)
         ]
 
     @console_ns.doc("create_api_based_extension")
@@ -130,9 +130,9 @@ class APIBasedExtensionAPI(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    def post(self):
+    @with_current_tenant_id
+    def post(self, current_tenant_id: str):
         payload = APIBasedExtensionPayload.model_validate(console_ns.payload or {})
-        _, current_tenant_id = current_account_with_tenant()
 
         extension_data = APIBasedExtension(
             tenant_id=current_tenant_id,
@@ -153,12 +153,12 @@ class APIBasedExtensionDetailAPI(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    def get(self, id: UUID):
+    @with_current_tenant_id
+    def get(self, current_tenant_id: str, id: UUID):
         api_based_extension_id = str(id)
-        _, tenant_id = current_account_with_tenant()
 
         return _serialize_api_based_extension(
-            APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
+            APIBasedExtensionService.get_with_tenant_id(current_tenant_id, api_based_extension_id)
         )
 
     @console_ns.doc("update_api_based_extension")
@@ -169,9 +169,9 @@ class APIBasedExtensionDetailAPI(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    def post(self, id: UUID):
+    @with_current_tenant_id
+    def post(self, current_tenant_id: str, id: UUID):
         api_based_extension_id = str(id)
-        _, current_tenant_id = current_account_with_tenant()
 
         extension_data_from_db = APIBasedExtensionService.get_with_tenant_id(current_tenant_id, api_based_extension_id)
 
@@ -197,9 +197,9 @@ class APIBasedExtensionDetailAPI(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    def delete(self, id: UUID):
+    @with_current_tenant_id
+    def delete(self, current_tenant_id: str, id: UUID):
         api_based_extension_id = str(id)
-        _, current_tenant_id = current_account_with_tenant()
 
         extension_data_from_db = APIBasedExtensionService.get_with_tenant_id(current_tenant_id, api_based_extension_id)
 
diff --git a/api/tests/unit_tests/controllers/console/test_extension.py b/api/tests/unit_tests/controllers/console/test_extension.py
index 20fc62073b..487cf8f54f 100644
--- a/api/tests/unit_tests/controllers/console/test_extension.py
+++ b/api/tests/unit_tests/controllers/console/test_extension.py
@@ -54,7 +54,6 @@ def _masked_api_key(api_key: str) -> str:
 def _mock_console_guards(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
     """Bypass console decorators so handlers can run in isolation."""
 
-    import controllers.console.extension as extension_module
     from controllers.console import wraps as wraps_module
 
     account = MagicMock()
@@ -66,7 +65,6 @@ def _mock_console_guards(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
     monkeypatch.setattr(wraps_module.dify_config, "EDITION", "CLOUD")
     monkeypatch.setattr("libs.login.dify_config.LOGIN_DISABLED", True)
     monkeypatch.delenv("INIT_PASSWORD", raising=False)
-    monkeypatch.setattr(extension_module, "current_account_with_tenant", lambda: (account, "tenant-123"))
     monkeypatch.setattr(wraps_module, "current_account_with_tenant", lambda: (account, "tenant-123"))
 
     # The login_required decorator consults the shared LocalProxy in libs.login.