diff --git a/api/commands/rbac.py b/api/commands/rbac.py index be4993920ad..e6ca34afd31 100644 --- a/api/commands/rbac.py +++ b/api/commands/rbac.py @@ -212,9 +212,7 @@ def migrate_member_roles_to_rbac( account_id=owner_account_id, member_account_ids=[account_id for account_id, _ in batch], ) - current_roles_by_account_id = { - item.account_id: {str(role.id) for role in item.roles} for item in current_roles - } + current_roles_by_account_id = {item.account_id: {role.id for role in item.roles} for item in current_roles} replace_jobs: list[tuple[str, str]] = [] for member_account_id, legacy_role in batch: @@ -468,7 +466,9 @@ def migrate_dataset_permissions_to_rbac( "account_id": operator_account_id, "dataset_id": current_dataset_id, "target_account_id": member_account_id, - "payload": replace_user_access_policies_payload.model_dump(mode="json"), + "payload": replace_user_access_policies_payload.model_dump( + mode="json", exclude_unset=True + ), }, }, } diff --git a/api/controllers/console/app/app.py b/api/controllers/console/app/app.py index 4f0022b37b4..fb07207376d 100644 --- a/api/controllers/console/app/app.py +++ b/api/controllers/console/app/app.py @@ -42,6 +42,7 @@ from controllers.console.wraps import ( from core.ops.ops_trace_manager import OpsTraceManager from core.rag.entities import PreProcessingRule, Rule, Segmentation from core.rag.retrieval.retrieval_methods import RetrievalMethod +from core.rbac import RBACResourceWhitelistScope from core.trigger.constants import TRIGGER_NODE_TYPES from extensions.ext_database import db from fields.base import ResponseModel @@ -68,6 +69,7 @@ from services.entities.knowledge_entities.knowledge_entities import ( WeightVectorSetting, ) from services.feature_service import FeatureService +from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task ALLOW_CREATE_APP_MODES = ["chat", "agent-chat", "advanced-chat", "workflow", "completion"] @@ -645,6 +647,14 @@ class AppListApi(Resource): app_service = AppService() app = app_service.create_app(current_tenant_id, params, current_user, session=db.session()) + if dify_config.RBAC_ENABLED: + enterprise_rbac_service.RBACService.AppAccess.replace_whitelist( + tenant_id=str(current_tenant_id), + account_id=current_user.id, + app_id=str(app.id), + payload=enterprise_rbac_service.ReplaceMemberBindings(scope=RBACResourceWhitelistScope.ALL), + ) + initialize_created_app_rbac_access_task.delay(current_tenant_id, current_user.id, app.id) permission_keys_map = enterprise_rbac_service.RBACService.AppPermissions.batch_get( str(current_tenant_id), current_user.id, diff --git a/api/controllers/console/workspace/rbac.py b/api/controllers/console/workspace/rbac.py index 39ad12080f3..878f91fe168 100644 --- a/api/controllers/console/workspace/rbac.py +++ b/api/controllers/console/workspace/rbac.py @@ -18,6 +18,7 @@ from extensions.ext_database import db from libs.login import current_account_with_tenant, login_required from models import Account from services.enterprise import rbac_service as svc +from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task class _RBACRoleList(svc.Paginated[svc.RBACRole]): @@ -596,14 +597,15 @@ class RBACAppWhitelistApi(Resource): def put(self, app_id): tenant_id, account_id = _current_ids() request = _payload(_ResourceAccessScopeRequest) - return _dump( - svc.RBACService.AppAccess.replace_whitelist( - tenant_id, - account_id, - str(app_id), - svc.ReplaceMemberBindings(scope=request.scope.value), - ) + result = svc.RBACService.AppAccess.replace_whitelist( + tenant_id, + account_id, + str(app_id), + svc.ReplaceMemberBindings(scope=request.scope.value), ) + if dify_config.RBAC_ENABLED and request.scope is RBACResourceWhitelistScope.ALL: + initialize_created_app_rbac_access_task.delay(tenant_id, account_id, str(app_id)) + return _dump(result) @console_ns.route("/workspaces/current/rbac/apps//user-access-policies") @@ -630,8 +632,8 @@ class RBACAppUserAccessPolicyAssignmentApi(Resource): svc.RBACService.AppAccess.replace_user_access_policies( tenant_id, account_id, - str(app_id), - str(target_account_id), + app_id, + target_account_id, payload, ) ) diff --git a/api/docker/entrypoint.sh b/api/docker/entrypoint.sh index 14778f49d5e..67832da826a 100755 --- a/api/docker/entrypoint.sh +++ b/api/docker/entrypoint.sh @@ -35,10 +35,10 @@ if [[ "${MODE}" == "worker" ]]; then if [[ -z "${CELERY_QUEUES}" ]]; then if [[ "${EDITION}" == "CLOUD" ]]; then # Cloud edition: separate queues for dataset and trigger tasks - DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution" + DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow_professional,workflow_team,workflow_sandbox,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution" else # Community edition (SELF_HOSTED): dataset, pipeline and workflow have separate queues - DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution" + DEFAULT_QUEUES="api_token,dataset,dataset_summary,priority_dataset,priority_pipeline,pipeline,mail,ops_trace,app_deletion,app_rbac,plugin,workflow_storage,conversation,workflow,schedule_poller,schedule_executor,triggered_workflow_dispatcher,trigger_refresh_publisher,trigger_refresh_executor,retention,workflow_based_app_execution" fi else DEFAULT_QUEUES="${CELERY_QUEUES}" diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 810f5f17bc7..9be77df955d 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -155,6 +155,7 @@ def init_app(app: DifyApp) -> Celery: "tasks.trigger_processing_tasks", # async trigger processing "tasks.generate_summary_index_task", # summary index generation "tasks.regenerate_summary_index_task", # summary index regeneration + "tasks.initialize_created_app_rbac_access_task", # app access initialization "tasks.app_generate.resume_agent_app_task", # ENG-635: Agent v2 chat ask_human resume ] day = dify_config.CELERY_BEAT_SCHEDULER_TIME diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index 556e0045b8d..12a02b6d3ad 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -20722,6 +20722,7 @@ Whitelist scopes accepted by RBAC app and dataset access config APIs. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | access_policy_ids | [ string ] | | No | +| account_ids | [ string ] | | No | #### ReplaceUserAccessPoliciesResponse diff --git a/api/services/account_service.py b/api/services/account_service.py index b5439467a23..99a06066ca6 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -10,6 +10,7 @@ import json import logging import secrets import uuid +from collections.abc import Iterator from datetime import UTC, datetime, timedelta from hashlib import sha256 from typing import Any, NotRequired, TypedDict, cast @@ -1563,6 +1564,25 @@ class TenantService: return updated_accounts + @staticmethod + def iter_member_account_id_batches(tenant_id: str, batch_size: int, *, session: Session) -> Iterator[list[str]]: + """Yield workspace member account ids in bounded, ordered batches.""" + offset = 0 + while True: + stmt = ( + select(TenantAccountJoin.account_id) + .where(TenantAccountJoin.tenant_id == tenant_id) + .order_by(TenantAccountJoin.id) + .offset(offset) + .limit(batch_size) + ) + account_ids = list(session.scalars(stmt).all()) + if not account_ids: + return + + yield account_ids + offset += batch_size + @staticmethod def get_dataset_operator_members(tenant: Tenant, *, session: Session) -> list[Account]: """Get dataset admin members""" diff --git a/api/services/enterprise/rbac_service.py b/api/services/enterprise/rbac_service.py index 47ac8d5aeae..b90f2ab6183 100644 --- a/api/services/enterprise/rbac_service.py +++ b/api/services/enterprise/rbac_service.py @@ -255,6 +255,7 @@ class ResourceUserAccessPoliciesResponse(_RBACModel): class ReplaceUserAccessPolicies(_RBACModel): access_policy_ids: list[str] = Field(default_factory=list) + account_ids: list[str] = Field(default_factory=list) @field_validator("access_policy_ids", mode="before") @classmethod @@ -1126,16 +1127,17 @@ class RBACService: tenant_id: str, account_id: str | None, app_id: str, - target_account_id: str, + target_account_id: str | None, payload: ReplaceUserAccessPolicies, ) -> ReplaceUserAccessPoliciesResponse: + request_data = payload.model_dump(mode="json") data = _inner_call( "PUT", f"{_INNER_PREFIX}/apps/user-access-policies", tenant_id=tenant_id, account_id=account_id, params={"app_id": app_id, "account_id": target_account_id}, - json=payload.model_dump(mode="json"), + json=request_data, ) return ReplaceUserAccessPoliciesResponse.model_validate(data or {}) @@ -1304,7 +1306,7 @@ class RBACService: tenant_id=tenant_id, account_id=account_id, params={"dataset_id": dataset_id, "account_id": target_account_id}, - json=payload.model_dump(mode="json"), + json=payload.model_dump(mode="json", exclude_unset=True), ) return ReplaceUserAccessPoliciesResponse.model_validate(data or {}) diff --git a/api/tasks/initialize_created_app_rbac_access_task.py b/api/tasks/initialize_created_app_rbac_access_task.py new file mode 100644 index 00000000000..c55126102e9 --- /dev/null +++ b/api/tasks/initialize_created_app_rbac_access_task.py @@ -0,0 +1,53 @@ +"""Initialize default RBAC access for existing workspace members after app creation.""" + +import logging + +from celery import shared_task + +from configs import dify_config +from extensions.ext_database import db +from services.account_service import TenantService +from services.enterprise import rbac_service as enterprise_rbac_service + +logger = logging.getLogger(__name__) + +APP_RBAC_ACCOUNT_POLICY_BATCH_SIZE = 500 +APP_RBAC_DEFAULT_ACCESS_POLICY_ID = "default" +APP_RBAC_QUEUE = "app_rbac" + + +@shared_task(queue=APP_RBAC_QUEUE, bind=True, max_retries=3, default_retry_delay=60) +def initialize_created_app_rbac_access_task(self, tenant_id: str, account_id: str, app_id: str) -> None: + """Grant the default app policy to current workspace members. + + App scope is persisted synchronously before this task is queued. Replacing + member policies is idempotent, so retrying the whole synchronization is safe + when the enterprise RBAC service is temporarily unavailable. + """ + if not dify_config.RBAC_ENABLED: + return + + try: + for account_ids in TenantService.iter_member_account_id_batches( + tenant_id, + APP_RBAC_ACCOUNT_POLICY_BATCH_SIZE, + session=db.session(), + ): + enterprise_rbac_service.RBACService.AppAccess.replace_user_access_policies( + tenant_id=tenant_id, + account_id=account_id, + app_id=app_id, + target_account_id=None, + payload=enterprise_rbac_service.ReplaceUserAccessPolicies( + access_policy_ids=[APP_RBAC_DEFAULT_ACCESS_POLICY_ID], + account_ids=account_ids, + ), + ) + except Exception as exc: + logger.exception( + "Failed to initialize app RBAC access; retrying: tenant_id=%s app_id=%s attempt=%s", + tenant_id, + app_id, + self.request.retries + 1, + ) + raise self.retry(exc=exc) diff --git a/api/tests/unit_tests/controllers/console/app/test_app_response_models.py b/api/tests/unit_tests/controllers/console/app/test_app_response_models.py index 79f109fd383..0bd45e41935 100644 --- a/api/tests/unit_tests/controllers/console/app/test_app_response_models.py +++ b/api/tests/unit_tests/controllers/console/app/test_app_response_models.py @@ -556,6 +556,7 @@ def test_app_create_api_attaches_permission_keys(app, app_module): with app.test_request_context("/apps", method="POST", json={}): with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(app_module.dify_config, "RBAC_ENABLED", True) app_module.console_ns.payload = { "name": "Created App", "description": "Summary", @@ -571,11 +572,25 @@ def test_app_create_api_attaches_permission_keys(app, app_module): "batch_get", lambda tenant_id, account_id, app_ids, session: {"app-new": ["app.acl.view_layout", "app.acl.edit"]}, ) + initialize_rbac_task = MagicMock() + monkeypatch.setattr( + app_module, + "initialize_created_app_rbac_access_task", + initialize_rbac_task, + ) + replace_whitelist = MagicMock() + monkeypatch.setattr( + app_module.enterprise_rbac_service.RBACService.AppAccess, + "replace_whitelist", + replace_whitelist, + ) resp, status = method(app_module.AppListApi(), "tenant-1", SimpleNamespace(id="acct-1")) assert status == 201 assert resp["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"] + assert replace_whitelist.call_args.kwargs["payload"].scope is app_module.RBACResourceWhitelistScope.ALL + initialize_rbac_task.delay.assert_called_once_with("tenant-1", "acct-1", "app-new") def test_app_list_api_attaches_permission_keys(app, app_module): diff --git a/api/tests/unit_tests/controllers/console/workspace/test_rbac.py b/api/tests/unit_tests/controllers/console/workspace/test_rbac.py index 92e819f73bf..a118f803a91 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_rbac.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_rbac.py @@ -262,6 +262,25 @@ class TestPaginationMapping: class TestResourceAccessScopeBindings: + def test_app_whitelist_all_schedules_member_policy_sync(self, app): + with ( + app.test_request_context( + "/workspaces/current/rbac/apps/app-1/whitelist", + method="PUT", + json={"scope": "all"}, + ), + patch("controllers.console.workspace.rbac._current_ids", return_value=("tenant-1", "acct-actor")), + patch("controllers.console.workspace.rbac.dify_config.RBAC_ENABLED", True), + patch( + "controllers.console.workspace.rbac.svc.RBACService.AppAccess.replace_whitelist", + return_value=rbac_mod.svc.ResourceWhitelist(), + ), + patch("controllers.console.workspace.rbac.initialize_created_app_rbac_access_task") as mock_sync_task, + ): + inspect.unwrap(rbac_mod.RBACAppWhitelistApi.put)(rbac_mod.RBACAppWhitelistApi(), "app-1") + + mock_sync_task.delay.assert_called_once_with("tenant-1", "acct-actor", "app-1") + def test_app_user_access_policy_assignment_forwards_ids(self, app): with ( app.test_request_context( diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index c37800d60e2..a39bc138598 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -666,6 +666,31 @@ class TestTenantService: sqlite_session.add(tenant_account_join) return tenant_account_join + def test_iter_member_account_id_batches_uses_offset_limit(self): + class FakeScalarResult: + def __init__(self, items: list[str]) -> None: + self.items = items + + def all(self) -> list[str]: + return self.items + + offsets: list[int] = [] + + def scalars(stmt): + offsets.append(stmt._offset_clause.value) + if len(offsets) == 1: + return FakeScalarResult(["acct-1", "acct-2"]) + if len(offsets) == 2: + return FakeScalarResult(["acct-3"]) + return FakeScalarResult([]) + + batches = list( + TenantService.iter_member_account_id_batches("tenant-1", 2, session=SimpleNamespace(scalars=scalars)) + ) + + assert batches == [["acct-1", "acct-2"], ["acct-3"]] + assert offsets == [0, 2, 4] + # ==================== get_account_role_in_tenant Tests ==================== # Backs the auth pipeline's `load_workspace_role`: None => non-member # (pipeline maps to 404), otherwise the caller's role (out-of-set role => 403). diff --git a/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py new file mode 100644 index 00000000000..2bc5b8a69aa --- /dev/null +++ b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py @@ -0,0 +1,69 @@ +from unittest.mock import MagicMock + +import pytest + +APP_RBAC_QUEUE = "app_rbac" + + +def test_initialize_created_app_rbac_access_task_uses_rbac_queue(): + from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task + + assert initialize_created_app_rbac_access_task.queue == APP_RBAC_QUEUE + + +def test_initialize_created_app_rbac_access_task_batches_workspace_members(monkeypatch): + import tasks.initialize_created_app_rbac_access_task as task_module + from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task + + monkeypatch.setattr(task_module.dify_config, "RBAC_ENABLED", True) + monkeypatch.setattr( + task_module.TenantService, + "iter_member_account_id_batches", + lambda tenant_id, batch_size, session: iter([["acct-1", "acct-2"], ["acct-3"]]), + ) + replace_whitelist = MagicMock() + replace_user_access_policies = MagicMock() + monkeypatch.setattr( + task_module.enterprise_rbac_service.RBACService.AppAccess, + "replace_whitelist", + replace_whitelist, + ) + monkeypatch.setattr( + task_module.enterprise_rbac_service.RBACService.AppAccess, + "replace_user_access_policies", + replace_user_access_policies, + ) + + initialize_created_app_rbac_access_task.run("tenant-1", "actor-1", "app-1") + + replace_whitelist.assert_not_called() + assert replace_user_access_policies.call_count == 2 + assert replace_user_access_policies.call_args_list[0].kwargs["payload"].account_ids == ["acct-1", "acct-2"] + assert replace_user_access_policies.call_args_list[1].kwargs["payload"].account_ids == ["acct-3"] + for call in replace_user_access_policies.call_args_list: + assert call.kwargs["payload"].access_policy_ids == [task_module.APP_RBAC_DEFAULT_ACCESS_POLICY_ID] + + +def test_initialize_created_app_rbac_access_task_retries_on_failure(monkeypatch): + import tasks.initialize_created_app_rbac_access_task as task_module + from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task + + monkeypatch.setattr(task_module.dify_config, "RBAC_ENABLED", True) + monkeypatch.setattr( + task_module.TenantService, + "iter_member_account_id_batches", + lambda tenant_id, batch_size, session: iter([["acct-1"]]), + ) + monkeypatch.setattr( + task_module.enterprise_rbac_service.RBACService.AppAccess, + "replace_user_access_policies", + MagicMock(side_effect=ConnectionError("RBAC unavailable")), + ) + retry = MagicMock(return_value=RuntimeError("retry requested")) + monkeypatch.setattr(initialize_created_app_rbac_access_task, "retry", retry) + + with pytest.raises(RuntimeError, match="retry requested"): + initialize_created_app_rbac_access_task.run("tenant-1", "actor-1", "app-1") + + retry.assert_called_once() + assert isinstance(retry.call_args.kwargs["exc"], ConnectionError) diff --git a/packages/contracts/generated/api/console/workspaces/types.gen.ts b/packages/contracts/generated/api/console/workspaces/types.gen.ts index 55792eccc56..f1874d56ff8 100644 --- a/packages/contracts/generated/api/console/workspaces/types.gen.ts +++ b/packages/contracts/generated/api/console/workspaces/types.gen.ts @@ -540,6 +540,7 @@ export type ResourceUserAccessPoliciesResponse = { export type ReplaceUserAccessPolicies = { access_policy_ids?: Array + account_ids?: Array } export type ReplaceUserAccessPoliciesResponse = { diff --git a/packages/contracts/generated/api/console/workspaces/zod.gen.ts b/packages/contracts/generated/api/console/workspaces/zod.gen.ts index 51fbbf5a6c9..a8a45e01c33 100644 --- a/packages/contracts/generated/api/console/workspaces/zod.gen.ts +++ b/packages/contracts/generated/api/console/workspaces/zod.gen.ts @@ -370,6 +370,7 @@ export const zDeleteMemberBindingsRequest = z.object({ */ export const zReplaceUserAccessPolicies = z.object({ access_policy_ids: z.array(z.string()).optional(), + account_ids: z.array(z.string()).optional(), }) /**