diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b18d16c6191..070880ffc24 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -267,5 +267,9 @@ /web/app/auth/refresh/route.ts @iamjoel @lyzno1 /web/service/server.ts @iamjoel @lyzno1 +# Frontend - Browser Mode Tests +/web/app/**/*.browser.spec.ts @lyzno1 +/web/app/**/*.browser.spec.tsx @lyzno1 + # Docker /docker/* @laipz8200 diff --git a/.github/workflows/accessibility-e2e.yml b/.github/workflows/accessibility-e2e.yml index 1840e0676b6..527a4fc8ffa 100644 --- a/.github/workflows/accessibility-e2e.yml +++ b/.github/workflows/accessibility-e2e.yml @@ -36,7 +36,7 @@ concurrency: jobs: accessibility: name: WCAG Level ${{ inputs.level == 'a' && 'A' || 'AA' }} · ${{ inputs.page }} - runs-on: depot-ubuntu-24.04-4 + runs-on: ubuntu-24.04 timeout-minutes: 120 defaults: run: diff --git a/.github/workflows/web-e2e.yml b/.github/workflows/web-e2e.yml index 3da6679a0b4..e6bc08d44af 100644 --- a/.github/workflows/web-e2e.yml +++ b/.github/workflows/web-e2e.yml @@ -14,7 +14,7 @@ permissions: jobs: test: name: Web Full-Stack E2E - runs-on: depot-ubuntu-24.04-4 + runs-on: ubuntu-24.04 timeout-minutes: 120 defaults: run: diff --git a/.github/workflows/web-tests.yml b/.github/workflows/web-tests.yml index 26687f64dff..6f098bffdfc 100644 --- a/.github/workflows/web-tests.yml +++ b/.github/workflows/web-tests.yml @@ -33,7 +33,7 @@ jobs: uses: ./.github/actions/setup-web - name: Run tests - run: vp test run --reporter=blob --reporter=minimal --shard=${{ matrix.shardIndex }}/${{ matrix.shardTotal }} --coverage + run: vp test run --project unit --reporter=blob --reporter=minimal --shard=${{ matrix.shardIndex }}/${{ matrix.shardTotal }} --coverage - name: Upload blob report if: ${{ !cancelled() }} @@ -46,7 +46,7 @@ jobs: browser-test: name: Web Browser Tests - runs-on: depot-ubuntu-24.04-4 + runs-on: ubuntu-24.04 timeout-minutes: 20 defaults: run: @@ -67,7 +67,17 @@ jobs: run: vp exec playwright install --with-deps --only-shell chromium - name: Run browser tests - run: vp test run --config vitest.browser.config.ts --silent=passed-only + run: vp test run --project browser --silent=passed-only + + - name: Upload browser failure artifacts + if: ${{ failure() && !cancelled() }} + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: browser-test-failures + path: web/.vitest-browser/ + if-no-files-found: ignore + include-hidden-files: true + retention-days: 7 merge-reports: name: Merge Test Reports @@ -111,7 +121,7 @@ jobs: dify-ui-test: name: dify-ui Tests - runs-on: depot-ubuntu-24.04-4 + runs-on: ubuntu-24.04 timeout-minutes: 20 env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} @@ -136,6 +146,16 @@ jobs: - name: Run dify-ui tests run: vp test run --project unit --coverage --silent=passed-only + - name: Upload dify-ui test failure artifacts + if: ${{ failure() && !cancelled() }} + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: dify-ui-test-failures + path: packages/dify-ui/.vitest-browser/ + if-no-files-found: ignore + include-hidden-files: true + retention-days: 7 + - name: Report coverage if: ${{ env.CODECOV_TOKEN != '' }} uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 @@ -147,7 +167,7 @@ jobs: dify-ui-storybook-test: name: dify-ui Storybook Tests - runs-on: depot-ubuntu-24.04-4 + runs-on: ubuntu-24.04 timeout-minutes: 20 defaults: run: @@ -169,3 +189,13 @@ jobs: - name: Run dify-ui Storybook tests run: vp run test:storybook + + - name: Upload dify-ui Storybook test failure artifacts + if: ${{ failure() && !cancelled() }} + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: dify-ui-storybook-test-failures + path: packages/dify-ui/.vitest-browser/ + if-no-files-found: ignore + include-hidden-files: true + retention-days: 7 diff --git a/.gitignore b/.gitignore index ee9885348f3..9449c7eee08 100644 --- a/.gitignore +++ b/.gitignore @@ -268,6 +268,7 @@ scripts/stress-test/reports/ .qoder/* .context/ # Vitest local reports +web/.vitest-browser/ web/.vitest-reports/ # dify-agent-runtime diff --git a/api/.env.example b/api/.env.example index a6326efcd96..a536800f2e8 100644 --- a/api/.env.example +++ b/api/.env.example @@ -709,6 +709,8 @@ AGENT_BACKEND_BASE_URL=http://localhost:5050 AGENT_BACKEND_API_TOKEN=dify-agent-run-token-for-dev-only AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS=30 AGENT_BACKEND_STREAM_MAX_RECONNECTS=3 +# Client deadline for converting a Binding file to a ToolFile through the Agent backend. +AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS=240 # KnowledgeFS (Dataset 2.0) KNOWLEDGE_FS_ENABLED=false diff --git a/api/.importlinter b/api/.importlinter index d7492912e7f..2de7120977d 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -93,6 +93,24 @@ forbidden_modules = sqlalchemy werkzeug +[importlinter:contract:account-application-boundary] +name = Account application services and contracts are framework and persistence neutral +type = forbidden +source_modules = + services.account_errors + services.account_ports + services.account_profile_service + services.entities.account_entities +forbidden_modules = + configs + controllers + extensions + flask + models + repositories + sqlalchemy + werkzeug + [importlinter:contract:app-definition-query-service-boundary] name = App definition query application service is framework and persistence neutral type = forbidden @@ -124,6 +142,24 @@ forbidden_modules = sqlalchemy werkzeug +[importlinter:contract:web-app-runtime-query-service-boundary] +name = Web app runtime query application service does not directly depend on configuration, feature implementation, transport, or ORM modules +type = forbidden +source_modules = + services.web_app_runtime_query_service +forbidden_modules = + configs + controllers + extensions + flask + models + repositories + services.feature_service + services.feature_service_gateway + sqlalchemy + werkzeug +allow_indirect_imports = True + [importlinter:contract:feature-query-service-boundary] name = Feature query application service is framework and persistence neutral type = forbidden @@ -200,3 +236,43 @@ forbidden_modules = repositories sqlalchemy werkzeug + +[importlinter:contract:tag-application-service-boundary] +name = Tag application service is framework and persistence neutral +type = forbidden +source_modules = + services.tag_application_service +forbidden_modules = + configs + controllers + extensions + flask + models + repositories + sqlalchemy + werkzeug + +[importlinter:contract:recommended-app-query-service-boundary] +name = Recommended app query application service is framework and persistence neutral +type = forbidden +source_modules = + services.recommended_app_query_service +forbidden_modules = + configs + controllers + extensions + flask + models + repositories + services.feature_service + services.recommended_app_catalog_gateway + sqlalchemy + werkzeug + +[importlinter:contract:recommended-app-catalog-gateway-boundary] +name = Recommended app catalog gateway does not depend on Flask +type = forbidden +source_modules = + services.recommended_app_catalog_gateway +forbidden_modules = + flask diff --git a/api/clients/agent_backend/factory.py b/api/clients/agent_backend/factory.py index d5af6aed486..2f4817dfcdd 100644 --- a/api/clients/agent_backend/factory.py +++ b/api/clients/agent_backend/factory.py @@ -8,10 +8,21 @@ from clients.agent_backend.client import AgentBackendRunClient, DifyAgentBackend from clients.agent_backend.fake_client import FakeAgentBackendRunClient, FakeAgentBackendScenario -def create_agent_backend_client(*, base_url: str, api_token: str | None = None, stream_timeout: float = 30) -> Client: +def create_agent_backend_client( + *, + base_url: str, + api_token: str | None = None, + stream_timeout: float = 30, + binding_file_download_timeout: float = 240, +) -> Client: api_token = api_token.strip() if api_token else None headers = {"Authorization": f"Bearer {api_token}"} if api_token else None - return Client(base_url=base_url, stream_timeout=stream_timeout, headers=headers) + return Client( + base_url=base_url, + stream_timeout=stream_timeout, + binding_file_download_timeout=binding_file_download_timeout, + headers=headers, + ) def create_agent_backend_run_client( diff --git a/api/configs/extra/__init__.py b/api/configs/extra/__init__.py index a142dbd7988..bba588e337b 100644 --- a/api/configs/extra/__init__.py +++ b/api/configs/extra/__init__.py @@ -1,6 +1,7 @@ from configs.extra.agent_backend_config import AgentBackendConfig from configs.extra.archive_config import ArchiveStorageConfig from configs.extra.knowledge_fs_config import KnowledgeFSConfig +from configs.extra.logstore_config import LogStoreConfig from configs.extra.notion_config import NotionConfig from configs.extra.sentry_config import SentryConfig from configs.extra.turnstile_config import TurnstileConfig @@ -11,6 +12,7 @@ class ExtraServiceConfig( AgentBackendConfig, ArchiveStorageConfig, KnowledgeFSConfig, + LogStoreConfig, NotionConfig, SentryConfig, TurnstileConfig, diff --git a/api/configs/extra/agent_backend_config.py b/api/configs/extra/agent_backend_config.py index 464b3e8ad17..2a43bc20159 100644 --- a/api/configs/extra/agent_backend_config.py +++ b/api/configs/extra/agent_backend_config.py @@ -37,6 +37,11 @@ class AgentBackendConfig(BaseSettings): default=3, ) + AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS: PositiveFloat = Field( + description="Client timeout for converting a Binding file to a ToolFile through the Agent backend.", + default=240, + ) + AGENT_SHELL_ENABLED: bool = Field( description=( "Inject the Home, Workspace, Sandbox, and Shell runtime layers into Agent runs. " diff --git a/api/configs/extra/logstore_config.py b/api/configs/extra/logstore_config.py new file mode 100644 index 00000000000..b9a9ed00758 --- /dev/null +++ b/api/configs/extra/logstore_config.py @@ -0,0 +1,11 @@ +from pydantic_settings import BaseSettings + + +class LogStoreConfig(BaseSettings): + """Migration controls for repositories backed by Aliyun LogStore.""" + + LOGSTORE_DUAL_WRITE_ENABLED: bool = False + + # Keep workflow graphs in LogStore by default. Deployments may disable this + # while migrating large graph payloads to another persistence owner. + LOGSTORE_ENABLE_PUT_GRAPH_FIELD: bool = True diff --git a/api/constants/__init__.py b/api/constants/__init__.py index 17220d7d15f..8f9e50a3b1d 100644 --- a/api/constants/__init__.py +++ b/api/constants/__init__.py @@ -36,6 +36,7 @@ _UNSTRUCTURED_DOCUMENT_EXTENSION_BASE: frozenset[str] = frozenset( "pptx", "xml", "epub", + "odt", ) ) _DEFAULT_DOCUMENT_EXTENSION_BASE: frozenset[str] = frozenset( @@ -53,6 +54,7 @@ _DEFAULT_DOCUMENT_EXTENSION_BASE: frozenset[str] = frozenset( "csv", "vtt", "properties", + "odt", ) ) diff --git a/api/constants/recommended_apps.json b/api/constants/recommended_apps.json index b7db9d64d06..6e9b6856542 100644 --- a/api/constants/recommended_apps.json +++ b/api/constants/recommended_apps.json @@ -166,6 +166,7 @@ "categories": ["Workflow"], "copyright": null, "description": "Based on users' choice, retrieve external knowledge to more accurately summarize articles.", + "is_learn_dify": true, "is_listed": true, "position": 5, "privacy_policy": null @@ -294,6 +295,7 @@ "categories": ["Workflow"], "copyright": null, "description": "Basic Workflow Template, a chatbot capable of identifying intents alongside with a knowledge base.", + "is_learn_dify": true, "is_listed": true, "position": 4, "privacy_policy": null @@ -326,6 +328,7 @@ "categories": ["Workflow"], "copyright": null, "description": "Basic Workflow Template, A chatbot with a knowledge base. ", + "is_learn_dify": true, "is_listed": true, "position": 4, "privacy_policy": null diff --git a/api/controllers/console/app/workflow.py b/api/controllers/console/app/workflow.py index d506a443a3a..4b9ff67c0e8 100644 --- a/api/controllers/console/app/workflow.py +++ b/api/controllers/console/app/workflow.py @@ -20,6 +20,8 @@ from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound import services +from configs import dify_config +from controllers.common.app_access import resolve_app_access_filter from controllers.common.controller_schemas import DefaultBlockConfigQuery, WorkflowListQuery, WorkflowUpdatePayload from controllers.common.errors import InvalidArgumentError from controllers.common.fields import GeneratedAppResponse, NewAppResponse, SimpleResultResponse @@ -53,6 +55,7 @@ from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.apps.workflow.app_generator import SKIP_PREPARE_USER_INPUTS_KEY from core.app.entities.app_invoke_entities import InvokeFrom from core.app.file_access import DatabaseFileAccessController +from core.db.session_factory import session_factory from core.helper import encrypter from core.helper.trace_id_helper import get_external_trace_id from core.plugin.impl.exc import PluginInvokeError @@ -95,7 +98,6 @@ from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, from services.errors.llm import InvokeRateLimitError from services.workflow_ref_service import WorkflowRefService from services.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError, WorkflowService -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -1303,7 +1305,7 @@ class PublishedWorkflowApi(Resource): workflow_service = WorkflowService() with sessionmaker(db.engine).begin() as session: - workflow, retirement_candidates = workflow_service.publish_workflow( + workflow = workflow_service.publish_workflow( session=session, app_model=app_model, account=current_user, @@ -1320,16 +1322,6 @@ class PublishedWorkflowApi(Resource): workflow_created_at = TimestampField().format(workflow.created_at) - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( - tenant_id=app_model.tenant_id, - agent_ids=retirement_candidates, - account_id=current_user.id, - ) - enqueue_agent_resource_collection( - tenant_id=app_model.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) return { "result": "success", "created_at": workflow_created_at, @@ -1620,10 +1612,11 @@ class WorkflowByIdApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) + @with_current_user @edit_permission_required @rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_EDIT) @console_ns.response(204, "Workflow deleted successfully") - def delete(self, app_model: App, workflow_id: str): + def delete(self, current_user: Account, app_model: App, workflow_id: str): """ Delete workflow """ @@ -1633,7 +1626,7 @@ class WorkflowByIdApi(Resource): # Create a session and manage the transaction with sessionmaker(db.engine).begin() as session: try: - workflow_service.delete_workflow( + retirement_candidates = workflow_service.delete_workflow( session=session, workflow_ref=workflow_ref, ) @@ -1644,6 +1637,11 @@ class WorkflowByIdApi(Resource): except ValueError as e: raise NotFound(str(e)) + WorkflowAgentRetirementService.retire_unowned( + tenant_id=app_model.tenant_id, + agent_ids=retirement_candidates, + account_id=current_user.id, + ) return None, 204 @@ -1929,8 +1927,9 @@ class WorkflowOnlineUsersApi(Resource): @setup_required @login_required @account_initialization_required + @with_current_user @with_current_tenant_id - def post(self, current_tenant_id: str): + def post(self, current_tenant_id: str, current_user: Account): args = WorkflowOnlineUsersPayload.model_validate(console_ns.payload or {}) app_ids = args.app_ids @@ -1940,8 +1939,20 @@ class WorkflowOnlineUsersApi(Resource): if not app_ids: return {"data": []} + access_filter = None workflow_service = WorkflowService() - accessible_app_ids = workflow_service.get_accessible_app_ids(app_ids, current_tenant_id, session=db.session()) + with session_factory.create_session() as session: + if dify_config.RBAC_ENABLED: + access_filter = resolve_app_access_filter(current_tenant_id, current_user.id, session=session) + app_maintainers = workflow_service.get_tenant_app_maintainers(app_ids, current_tenant_id, session=session) + + accessible_app_ids = set(app_maintainers) + if access_filter is not None: + accessible_app_ids = { + app_id + for app_id, maintainer in app_maintainers.items() + if access_filter.is_app_accessible(app_id, maintainer, current_user.id) + } ordered_accessible_app_ids = [app_id for app_id in app_ids if app_id in accessible_app_ids] users_json_by_app_id: dict[str, Any] = {} diff --git a/api/controllers/console/app/wraps.py b/api/controllers/console/app/wraps.py index 04047f86c8b..e844348de85 100644 --- a/api/controllers/console/app/wraps.py +++ b/api/controllers/console/app/wraps.py @@ -1,8 +1,9 @@ """Controller decorators for console app resources. `get_app_model` still supports legacy handlers backed by Flask-SQLAlchemy's -scoped session. Trial app handlers compose `get_app_model_with_trial` under -`controllers.common.session.with_session` and always reuse that request session. +scoped session. Preview handlers compose `get_previewable_app_model` under +`controllers.common.session.with_session`; preview admission finishes before +the request Session loads the accepted App. """ from collections.abc import Callable @@ -16,16 +17,17 @@ from configs import dify_config from controllers.common.session import with_session from controllers.common.wraps import RBACPermission, RBACResourceScope, enforce_rbac_access from controllers.console.app.error import AppNotFoundError +from extensions.ext_application_services import application_services from extensions.ext_database import db from libs.login import current_account_with_tenant -from models import App, AppMode, TrialApp +from models import App, AppMode from models.agent import AgentScope -from services.recommended_app_service import RecommendedAppService +from services.app_service import AppService __all__ = [ "agent_manage_required_for_agent_app", "get_app_model", - "get_app_model_with_trial", + "get_previewable_app_model", "with_session", ] @@ -48,12 +50,11 @@ def _load_app_model_from_scoped_session(app_id: str) -> App | None: return app_model -def _load_app_model_with_trial(session: Session, app_id: str) -> App | None: - """Load a normal app through its trial registration without applying current-tenant scope.""" - app_model = session.scalar( - select(App).join(TrialApp, TrialApp.app_id == App.id).where(App.id == app_id, App.status == "normal").limit(1) - ) - return app_model +def _load_previewable_app_model(session: Session, app_id: str) -> App | None: + """Load a normal App after preview admission completes outside the request Session.""" + if not application_services().recommended_app_queries.is_previewable(app_id): + return None + return AppService.get_normal_app_by_id(app_id, session) def agent_manage_required_for_agent_app[**P, R](view: Callable[P, R]) -> Callable[P, R]: @@ -183,7 +184,7 @@ def get_app_model[**P, R]( @overload -def get_app_model_with_trial[**P, R]( +def get_previewable_app_model[**P, R]( view: Callable[P, R], *, mode: AppMode | list[AppMode] | None = None, @@ -191,19 +192,25 @@ def get_app_model_with_trial[**P, R]( @overload -def get_app_model_with_trial[**P, R]( +def get_previewable_app_model[**P, R]( view: None = None, *, mode: AppMode | list[AppMode] | None = None, ) -> Callable[[Callable[P, R]], Callable[P, R]]: ... -def get_app_model_with_trial[**P, R]( +def get_previewable_app_model[**P, R]( view: Callable[P, R] | None = None, *, mode: AppMode | list[AppMode] | None = None, ) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]: - """Inject a trial-registered or recommended App using the Session supplied by `with_session`.""" + """Inject an App authorized for read-only template preview. + + Preview reads accept either an explicit TrialApp registration or membership + in the recommended catalog. This does not grant trial execution, which is + separately protected by TrialAppResource's feature, registration, and quota + checks. + """ def decorator(view_func: Callable[P, R]) -> Callable[P, R]: @wraps(view_func) @@ -218,10 +225,8 @@ def get_app_model_with_trial[**P, R]( session = _get_injected_session(args) if session is None: - raise RuntimeError("get_app_model_with_trial requires @with_session") - app_model = _load_app_model_with_trial(session, app_id) - if app_model is None: - app_model = RecommendedAppService.get_app(app_id, session=session) + raise RuntimeError("get_previewable_app_model requires @with_session") + app_model = _load_previewable_app_model(session, app_id) if not app_model: raise AppNotFoundError() diff --git a/api/controllers/console/auth/activate.py b/api/controllers/console/auth/activate.py index 3921414369e..69ed1899f2e 100644 --- a/api/controllers/console/auth/activate.py +++ b/api/controllers/console/auth/activate.py @@ -6,11 +6,14 @@ from constants.languages import supported_language from controllers.common.schema import query_params_from_model, register_schema_models from controllers.console import console_ns from controllers.console.auth.error import InvitationAccountMismatchError as InvitationAccountMismatchHTTPError -from controllers.console.error import AccountInFreezeError, AlreadyActivateError +from controllers.console.error import AccountInFreezeError, AlreadyActivateError, EmailDomainSuspendedError from extensions.ext_application_services import application_services from libs.helper import EmailStr, dump_response, timezone from libs.login import current_account_with_tenant from libs.token import extract_access_token +from services.account_activation_service import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) from services.account_activation_service import ( FrozenAccountError, InvalidInvitationError, @@ -141,6 +144,8 @@ class ActivateApi(Resource): raise AlreadyActivateError() from None except InvitationAccountMismatchError: raise InvitationAccountMismatchHTTPError() from None + except EmailDomainSuspendedRegistrationError: + raise EmailDomainSuspendedError() from None except FrozenAccountError: raise AccountInFreezeError() from None diff --git a/api/controllers/console/auth/email_register.py b/api/controllers/console/auth/email_register.py index c335fb852f4..8d5244b5312 100644 --- a/api/controllers/console/auth/email_register.py +++ b/api/controllers/console/auth/email_register.py @@ -24,9 +24,15 @@ from libs.password import valid_password from models import Account from services.account_service import AccountService from services.billing_service import BillingService -from services.errors.account import AccountRegisterError, SeatsLimitExceededError +from services.errors.account import ( + AccountRegisterError, + SeatsLimitExceededError, +) +from services.errors.account import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) -from ..error import AccountInFreezeError, EmailSendIpLimitError, SeatsLimitExceeded +from ..error import AccountInFreezeError, EmailDomainSuspendedError, EmailSendIpLimitError, SeatsLimitExceeded from ..wraps import email_password_login_enabled, email_register_enabled, model_validate, setup_required @@ -99,10 +105,12 @@ class EmailRegisterSendEmailApi(Resource): if req_data.language is not None and req_data.language in languages: language = req_data.language - if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and BillingService.is_email_in_freeze( - normalized_email - ): - raise AccountInFreezeError() + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + freeze_type = BillingService.get_email_freeze_type(normalized_email) + if freeze_type: + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedError() + raise AccountInFreezeError() account = AccountService.get_account_by_email_with_case_fallback(req_data.email, session=db.session()) token = AccountService.send_email_register_email(email=normalized_email, account=account, language=language) @@ -217,5 +225,7 @@ class EmailRegisterResetApi(Resource): ) except SeatsLimitExceededError: raise SeatsLimitExceeded() - except AccountRegisterError: - raise AccountInFreezeError() + except EmailDomainSuspendedRegistrationError as exc: + raise EmailDomainSuspendedError() from exc + except AccountRegisterError as exc: + raise AccountInFreezeError() from exc diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 6629784afaf..4252b4373c3 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -34,6 +34,7 @@ from controllers.console.error import ( AccountBannedError, AccountInFreezeError, AccountNotFound, + EmailDomainSuspendedError, EmailSendIpLimitError, NotAllowedCreateWorkspace, SeatsLimitExceeded, @@ -74,6 +75,9 @@ from services.errors.account import ( RefreshTokenNotFoundError, SeatsLimitExceededError, ) +from services.errors.account import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError from services.feature_service import FeatureService from services.turnstile_service import ( @@ -149,11 +153,13 @@ class LoginApi(Resource): request_email = req_data.email normalized_email = request_email.lower() - if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and BillingService.is_email_in_freeze( - normalized_email - ): - _log_console_login_failure(email=normalized_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE) - raise AccountInFreezeError() + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + freeze_type = BillingService.get_email_freeze_type(normalized_email) + if freeze_type: + _log_console_login_failure(email=normalized_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE) + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedError() + raise AccountInFreezeError() is_login_error_rate_limit = AccountService.is_login_error_rate_limit(normalized_email) if is_login_error_rate_limit: @@ -255,8 +261,10 @@ class ResetPasswordSendEmailApi(Resource): language = "en-US" try: account = _get_account_with_case_fallback(req_data.email) - except AccountRegisterError: - raise AccountInFreezeError() + except EmailDomainSuspendedRegistrationError as exc: + raise EmailDomainSuspendedError() from exc + except AccountRegisterError as exc: + raise AccountInFreezeError() from exc token = AccountService.send_reset_password_email( email=normalized_email, @@ -297,8 +305,10 @@ class EmailCodeLoginSendEmailApi(Resource): language = "en-US" try: account = _get_account_with_case_fallback(req_data.email) - except AccountRegisterError: - raise AccountInFreezeError() + except EmailDomainSuspendedRegistrationError as exc: + raise EmailDomainSuspendedError() from exc + except AccountRegisterError as exc: + raise AccountInFreezeError() from exc if account is None: if FeatureService.get_system_features().is_allow_register: @@ -377,9 +387,12 @@ class EmailCodeLoginApi(Resource): except Unauthorized as exc: _log_console_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_BANNED) raise AccountBannedError() from exc - except AccountRegisterError: + except EmailDomainSuspendedRegistrationError as exc: _log_console_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE) - raise AccountInFreezeError() + raise EmailDomainSuspendedError() from exc + except AccountRegisterError as exc: + _log_console_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE) + raise AccountInFreezeError() from exc if account: tenants = TenantService.get_join_tenants(account, session=db.session()) if not tenants: @@ -405,9 +418,12 @@ class EmailCodeLoginApi(Resource): raise NotAllowedCreateWorkspace() except SeatsLimitExceededError: raise SeatsLimitExceeded() - except AccountRegisterError: + except EmailDomainSuspendedRegistrationError as exc: _log_console_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE) - raise AccountInFreezeError() + raise EmailDomainSuspendedError() from exc + except AccountRegisterError as exc: + _log_console_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE) + raise AccountInFreezeError() from exc except WorkspacesLimitExceededError: raise WorkspacesLimitExceeded() token_pair = AccountService.login(account, session=db.session(), ip_address=ip_address) diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index 41b94431f9a..e7cc10c9e06 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -12,6 +12,7 @@ from configs import dify_config from constants.languages import languages from controllers.common.fields import RedirectResponse from controllers.common.schema import query_params_from_model, register_response_schema_model, register_schema_models +from controllers.console.error import AccountInFreezeError, EmailDomainSuspendedError from enums import DeploymentEdition from extensions.ext_database import db from libs.datetime_utils import naive_utc_now @@ -26,7 +27,14 @@ from libs.token import ( from models import Account, AccountStatus from services.account_service import AccountService, RegisterService, TenantService from services.billing_service import BillingService -from services.errors.account import AccountNotFoundError, AccountRegisterError, SeatsLimitExceededError +from services.errors.account import ( + AccountNotFoundError, + AccountRegisterError, + SeatsLimitExceededError, +) +from services.errors.account import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkSpaceNotFoundError from services.feature_service import FeatureService @@ -249,8 +257,10 @@ class OAuthCallback(Resource): ) except SeatsLimitExceededError: return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Licensed seats limit exceeded.") - except AccountRegisterError as e: - return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message={e.description}") + except EmailDomainSuspendedRegistrationError: + return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message={EmailDomainSuspendedError.description}") + except AccountRegisterError as exc: + return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message={exc.description}") # Check account status if account.status == AccountStatus.BANNED: @@ -309,15 +319,12 @@ def _generate_account( normalized_email = user_info.email.lower() oauth_new_user = True if not FeatureService.get_system_features().is_allow_register: - if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and BillingService.is_email_in_freeze( - normalized_email - ): - raise AccountRegisterError( - description=( - "This email account has been deleted within the past " - "30 days and is temporarily unavailable for new account registration" - ) - ) + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + freeze_type = BillingService.get_email_freeze_type(normalized_email) + if freeze_type: + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedRegistrationError() + raise AccountRegisterError(description=AccountInFreezeError.description or "") raise AccountRegisterError(description=("Invalid email or password")) account_name = user_info.name or "Dify" interface_language = _preferred_interface_language(language) diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 7dead9ba17a..56f33d9f393 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -732,10 +732,12 @@ class DatasetApi(Resource): dataset = DatasetService.get_dataset(dataset_id_str, session) if dataset is None: raise NotFound("Dataset not found.") - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) + + if not dify_config.RBAC_ENABLED: + try: + DatasetService.check_dataset_permission(dataset, current_user, session) + except services.errors.account.NoPermissionError as e: + raise Forbidden(str(e)) permissions = enterprise_rbac_service.RBACService.MyPermissions.get( current_tenant_id, current_user.id, diff --git a/api/controllers/console/datasets/datasets_segments.py b/api/controllers/console/datasets/datasets_segments.py index d1fa14902e6..35b2256a118 100644 --- a/api/controllers/console/datasets/datasets_segments.py +++ b/api/controllers/console/datasets/datasets_segments.py @@ -631,12 +631,16 @@ class DatasetDocumentSegmentBatchImportApi(Resource): ): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) + dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) if not dataset: raise NotFound("Dataset not found.") + # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) + document_ref = DatasetRefService.create_document_ref_from_id( + DatasetRefService.create_dataset_ref(dataset), document_id_str + ) + document = DatasetRefService.get_document_by_ref(document_ref, session=session) if not document: raise NotFound("Document not found.") diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py index a28b44d0346..95f19e11fd8 100644 --- a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py +++ b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py @@ -336,6 +336,8 @@ class DatasourceAuth(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required + @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.CREDENTIAL_MANAGE, resource_required=False) @with_current_user @with_current_tenant_id def get(self, current_tenant_id: str, user: Account, provider_id: str): diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py index 63a14083c85..7c5e0cff893 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py @@ -61,6 +61,7 @@ from models import Account from models.dataset import Pipeline from models.model import EndUser from models.workflow import Workflow +from services.agent.retirement_service import WorkflowAgentRetirementService from services.dataset_service import DatasetService from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.errors.llm import InvokeRateLimitError @@ -770,8 +771,9 @@ class RagPipelineByIdApi(Resource): @account_initialization_required @edit_permission_required @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) + @with_current_user @get_rag_pipeline - def delete(self, pipeline: Pipeline, workflow_id: str): + def delete(self, current_user: Account, pipeline: Pipeline, workflow_id: str): """ Delete a published workflow version that is not currently active on the pipeline. """ @@ -783,7 +785,7 @@ class RagPipelineByIdApi(Resource): with sessionmaker(db.engine).begin() as session: try: - workflow_service.delete_workflow( + retirement_candidates = workflow_service.delete_workflow( session=session, workflow_ref=workflow_ref, ) @@ -794,6 +796,11 @@ class RagPipelineByIdApi(Resource): except ValueError as e: raise NotFound(str(e)) + WorkflowAgentRetirementService.retire_unowned( + tenant_id=pipeline.tenant_id, + agent_ids=retirement_candidates, + account_id=current_user.id, + ) return None, 204 diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py index e4352f92f88..0ecb552bdb9 100644 --- a/api/controllers/console/error.py +++ b/api/controllers/console/error.py @@ -92,11 +92,17 @@ class AccountInFreezeError(BaseHTTPException): error_code = "account_in_freeze" code = 400 description = ( - "This email account has been deleted within the past 30 days" + "This email account has been deleted within the past 30 days " "and is temporarily unavailable for new account registration." ) +class EmailDomainSuspendedError(BaseHTTPException): + error_code = "email_domain_suspended" + code = 400 + description = "This email domain has been suspended." + + class EducationVerifyLimitError(BaseHTTPException): error_code = "education_verify_limit" description = "Rate limit exceeded" diff --git a/api/controllers/console/explore/error.py b/api/controllers/console/explore/error.py index e96fa64f846..6f0ba500b84 100644 --- a/api/controllers/console/explore/error.py +++ b/api/controllers/console/explore/error.py @@ -31,6 +31,12 @@ class AppAccessDeniedError(BaseHTTPException): code = 403 +class RecommendedAppNotFoundError(BaseHTTPException): + error_code = "recommended_app_not_found" + description = "Recommended app not found." + code = 404 + + class TrialAppNotAllowed(BaseHTTPException): """*403* `Trial App Not Allowed` @@ -51,3 +57,9 @@ class TrialAppLimitExceeded(BaseHTTPException): error_code = "trial_app_limit_exceeded" code = 403 description = "The user has exceeded the trial app limit." + + +class TrialAppFeatureDisabledError(BaseHTTPException): + error_code = "trial_app_feature_disabled" + code = 403 + description = "Trial app feature is not enabled." diff --git a/api/controllers/console/explore/recommended_app.py b/api/controllers/console/explore/recommended_app.py index aa91cd2d6c9..1125c0bc580 100644 --- a/api/controllers/console/explore/recommended_app.py +++ b/api/controllers/console/explore/recommended_app.py @@ -2,18 +2,18 @@ from typing import Any from uuid import UUID from flask_restx import Resource -from pydantic import BaseModel, Field, RootModel, computed_field, field_validator +from pydantic import BaseModel, Field, computed_field, field_validator -from constants.languages import languages from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from controllers.console import console_ns +from controllers.console.explore.error import RecommendedAppNotFoundError from controllers.console.wraps import account_initialization_required, model_validate, with_current_user -from extensions.ext_database import db +from extensions.ext_application_services import application_services from fields.base import ResponseModel from libs.helper import build_icon_url, dump_response from libs.login import login_required from models import Account -from services.recommended_app_service import RecommendedAppService +from services.recommended_app_query_service import RecommendedAppNotFoundError as RecommendedAppQueryNotFoundError class RecommendedAppsQuery(BaseModel): @@ -79,10 +79,6 @@ class RecommendedAppDetailResponse(ResponseModel): can_trial: bool -class RecommendedAppDetailNullableResponse(RootModel[RecommendedAppDetailResponse | None]): - pass - - register_schema_models( console_ns, RecommendedAppsQuery, @@ -94,18 +90,9 @@ register_response_schema_models( RecommendedAppListResponse, LearnDifyAppListResponse, RecommendedAppDetailResponse, - RecommendedAppDetailNullableResponse, ) -def _resolve_language(language: str | None, user: Account) -> str: - if language and language in languages: - return language - if user.interface_language: - return user.interface_language - return languages[0] - - @console_ns.route("/explore/apps") class RecommendedAppListApi(Resource): @console_ns.doc(params=query_params_from_model(RecommendedAppsQuery)) @@ -115,12 +102,12 @@ class RecommendedAppListApi(Resource): @with_current_user @model_validate(RecommendedAppsQuery) def get(self, req_data: RecommendedAppsQuery, current_user: Account): - # language args - language_prefix = _resolve_language(req_data.language, current_user) - return dump_response( RecommendedAppListResponse, - RecommendedAppService.get_recommended_apps_and_categories(language_prefix, session=db.session()), + application_services().recommended_app_queries.list_recommended( + requested_language=req_data.language, + interface_language=current_user.interface_language, + ), ) @@ -133,19 +120,24 @@ class LearnDifyAppListApi(Resource): @with_current_user @model_validate(RecommendedAppsQuery) def get(self, req_data: RecommendedAppsQuery, current_user: Account): - language_prefix = _resolve_language(req_data.language, current_user) - return dump_response( LearnDifyAppListResponse, - RecommendedAppService.get_learn_dify_apps(language_prefix, session=db.session()), + application_services().recommended_app_queries.list_learn_dify( + requested_language=req_data.language, + interface_language=current_user.interface_language, + ), ) @console_ns.route("/explore/apps/") class RecommendedAppApi(Resource): - @console_ns.response(200, "Success", console_ns.models[RecommendedAppDetailNullableResponse.__name__]) + @console_ns.response(200, "Success", console_ns.models[RecommendedAppDetailResponse.__name__]) + @console_ns.response(404, "Recommended app not found") @login_required @account_initialization_required def get(self, app_id: UUID): - result = RecommendedAppService.get_recommend_app_detail(str(app_id), session=db.session()) - return RecommendedAppDetailNullableResponse.model_validate(result).model_dump(mode="json") + try: + result = application_services().recommended_app_queries.get_detail(str(app_id)) + except RecommendedAppQueryNotFoundError: + raise RecommendedAppNotFoundError() from None + return dump_response(RecommendedAppDetailResponse, result) diff --git a/api/controllers/console/explore/trial.py b/api/controllers/console/explore/trial.py index bde115b6f25..d5388c7b542 100644 --- a/api/controllers/console/explore/trial.py +++ b/api/controllers/console/explore/trial.py @@ -38,7 +38,7 @@ from controllers.console.app.error import ( SpeechToTextDisabledError, UnsupportedAudioTypeError, ) -from controllers.console.app.wraps import get_app_model_with_trial, with_session +from controllers.console.app.wraps import get_previewable_app_model, with_session from controllers.console.explore.error import ( AppSuggestedQuestionsAfterAnswerDisabledError, NotChatAppError, @@ -96,7 +96,6 @@ from services.errors.message import ( SuggestedQuestionsAfterAnswerDisabledError, ) from services.message_service import MessageService -from services.recommended_app_service import RecommendedAppService logger = logging.getLogger(__name__) @@ -511,7 +510,7 @@ class TrialAppWorkflowRunApi(TrialAppResource): invoke_from=InvokeFrom.EXPLORE, streaming=True, ) - RecommendedAppService.add_trial_app_record(app_id, user_id, session=session) + application_services().trial_app_usage.record(app_id=app_id, account_id=user_id) # response-contract:ignore compact_generate_response return helper.compact_generate_response(response) except ProviderTokenNotInitError as ex: @@ -589,7 +588,7 @@ class TrialChatApi(TrialAppResource): invoke_from=InvokeFrom.EXPLORE, streaming=True, ) - RecommendedAppService.add_trial_app_record(app_id, user_id, session=session) + application_services().trial_app_usage.record(app_id=app_id, account_id=user_id) # response-contract:ignore compact_generate_response return helper.compact_generate_response(response) except services.errors.conversation.ConversationNotExistsError: @@ -675,7 +674,7 @@ class TrialChatAudioApi(TrialAppResource): session=db.session(), end_user=None, ) - RecommendedAppService.add_trial_app_record(app_id, user_id, session=db.session()) + application_services().trial_app_usage.record(app_id=app_id, account_id=user_id) return response except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") @@ -736,7 +735,7 @@ class TrialChatTextApi(TrialAppResource): voice=voice, message_ref=message_ref, ) - RecommendedAppService.add_trial_app_record(app_id, user_id, session=db.session()) + application_services().trial_app_usage.record(app_id=app_id, account_id=user_id) return response except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") @@ -794,7 +793,7 @@ class TrialCompletionApi(TrialAppResource): streaming=streaming, ) - RecommendedAppService.add_trial_app_record(app_id, user_id, session=session) + application_services().trial_app_usage.record(app_id=app_id, account_id=user_id) # response-contract:ignore compact_generate_response return helper.compact_generate_response(response) except services.errors.conversation.ConversationNotExistsError: @@ -824,7 +823,7 @@ class TrialSitApi(Resource): @console_ns.response(200, "Success", console_ns.models[SiteResponse.__name__]) @with_session(write=False) - @get_app_model_with_trial(None) + @get_previewable_app_model(None) def get(self, session: Session, app_model): """Retrieve app site info. @@ -848,7 +847,7 @@ class TrialAppParameterApi(Resource): @console_ns.response(200, "Success", console_ns.models[ParametersResponse.__name__]) @with_session(write=False) - @get_app_model_with_trial(None) + @get_previewable_app_model(None) def get(self, session: Session, app_model): """Retrieve app parameters.""" @@ -866,7 +865,7 @@ class TrialAppParameterApi(Resource): class AppApi(Resource): @console_ns.response(200, "Success", console_ns.models[TrialAppDetailResponse.__name__]) @with_session(write=False) - @get_app_model_with_trial(None) + @get_previewable_app_model(None) def get(self, session: Session, app_model): """Get app detail""" @@ -882,7 +881,7 @@ class AppApi(Resource): class AppWorkflowApi(Resource): @console_ns.response(200, "Success", console_ns.models[TrialWorkflowResponse.__name__]) @with_session(write=False) - @get_app_model_with_trial(None) + @get_previewable_app_model(None) def get(self, session: Session, app_model): """Get workflow detail""" if not app_model.workflow_id: @@ -902,7 +901,7 @@ class DatasetListApi(Resource): @console_ns.doc(params=query_params_from_model(TrialDatasetListQuery)) @console_ns.response(200, "Success", console_ns.models[TrialDatasetListResponse.__name__]) @with_session(write=False) - @get_app_model_with_trial(None) + @get_previewable_app_model(None) def get(self, session: Session, app_model): page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) diff --git a/api/controllers/console/explore/wraps.py b/api/controllers/console/explore/wraps.py index 1f4da57f9aa..a09341f649d 100644 --- a/api/controllers/console/explore/wraps.py +++ b/api/controllers/console/explore/wraps.py @@ -2,19 +2,23 @@ from collections.abc import Callable from functools import wraps from typing import Concatenate -from flask import abort from flask_restx import Resource from sqlalchemy import select from werkzeug.exceptions import NotFound -from controllers.console.explore.error import AppAccessDeniedError, TrialAppLimitExceeded, TrialAppNotAllowed +from controllers.console.explore.error import ( + AppAccessDeniedError, + TrialAppFeatureDisabledError, + TrialAppLimitExceeded, + TrialAppNotAllowed, +) from controllers.console.wraps import account_initialization_required +from extensions.ext_application_services import application_services from extensions.ext_database import db from libs.login import current_account_with_tenant, login_required from models import AccountTrialAppRecord, App, InstalledApp, TrialApp from services.enterprise.enterprise_service import EnterpriseService from services.feature_service import FeatureService -from services.recommended_app_service import RecommendedAppService def installed_app_required[**P, R](view: Callable[Concatenate[InstalledApp, P], R] | None = None): @@ -107,8 +111,8 @@ def trial_app_required[**P, R](view: Callable[Concatenate[App, P], R] | None = N def trial_feature_enable[**P, R](view: Callable[P, R]): @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): - if not RecommendedAppService.is_trial_app_enabled(): - abort(403, "Trial app feature is not enabled.") + if not application_services().recommended_app_queries.is_trial_enabled(): + raise TrialAppFeatureDisabledError() return view(*args, **kwargs) return decorated diff --git a/api/controllers/console/snippets/payloads.py b/api/controllers/console/snippets/payloads.py index de4115b6576..86b7735de6f 100644 --- a/api/controllers/console/snippets/payloads.py +++ b/api/controllers/console/snippets/payloads.py @@ -167,3 +167,9 @@ class IncludeSecretQuery(BaseModel): """Query parameter for including secret variables in export.""" include_secret: str = Field(default="false", description="Whether to include secret variables") + + +class SnippetExportQuery(IncludeSecretQuery): + """Query parameters for exporting a snippet workflow as DSL.""" + + workflow_id: str | None = Field(default=None, description="Specific published workflow version to export") diff --git a/api/controllers/console/snippets/snippet_workflow.py b/api/controllers/console/snippets/snippet_workflow.py index d1399801a06..309a9d9ea1a 100644 --- a/api/controllers/console/snippets/snippet_workflow.py +++ b/api/controllers/console/snippets/snippet_workflow.py @@ -57,12 +57,11 @@ from libs.helper import TimestampField from libs.login import current_account_with_tenant, login_required from models import Account from models.snippet import CustomizedSnippet -from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.workflow_publish_service import WorkflowAgentPublishService from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError +from services.errors.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError from services.snippet_generate_service import SnippetGenerateService from services.snippet_service import SnippetService -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -298,9 +297,8 @@ class SnippetPublishedWorkflowApi(Resource): with Session(db.engine) as session: snippet = session.merge(snippet) - tenant_id = snippet.tenant_id try: - workflow, retirement_candidates = snippet_service.publish_workflow( + workflow = snippet_service.publish_workflow( session=session, snippet=snippet, account=current_user, @@ -310,16 +308,6 @@ class SnippetPublishedWorkflowApi(Resource): except ValueError as e: return {"message": str(e)}, 400 - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( - tenant_id=tenant_id, - agent_ids=retirement_candidates, - account_id=current_user.id, - ) - enqueue_agent_resource_collection( - tenant_id=tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) return { "result": "success", "created_at": workflow_created_at, @@ -480,6 +468,39 @@ class SnippetWorkflowByIdApi(Resource): response["input_fields"] = snippet.input_fields_list return response + @console_ns.doc("delete_snippet_workflow_by_id") + @console_ns.doc(description="Delete a published snippet workflow version") + @console_ns.doc(params={"snippet_id": "Snippet ID", "workflow_id": "Workflow ID"}) + @console_ns.response(204, "Workflow deleted successfully") + @console_ns.response(400, "Workflow is in use") + @console_ns.response(404, "Workflow not found") + @setup_required + @login_required + @account_initialization_required + @get_snippet + @edit_permission_required + @rbac_permission_required( + RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False + ) + def delete(self, snippet: CustomizedSnippet, workflow_id: str): + """Delete a published snippet workflow version.""" + snippet_service = _snippet_service() + with _snippet_session_maker().begin() as session: + try: + snippet_service.delete_workflow( + session=session, + snippet=snippet, + workflow_id=workflow_id, + ) + except WorkflowInUseError as e: + raise BadRequest(str(e)) + except DraftWorkflowDeletionError as e: + raise BadRequest(str(e)) + except ValueError as e: + raise NotFound(str(e)) + + return None, 204 + @console_ns.route("/snippets//workflow-runs") class SnippetWorkflowRunsApi(Resource): diff --git a/api/controllers/console/tag/tags.py b/api/controllers/console/tag/tags.py index 0084e614c5b..14a3ee96535 100644 --- a/api/controllers/console/tag/tags.py +++ b/api/controllers/console/tag/tags.py @@ -3,37 +3,32 @@ from uuid import UUID from flask_restx import Resource from pydantic import BaseModel, Field, RootModel, field_validator -from sqlalchemy import select -from werkzeug.exceptions import Forbidden +from werkzeug.exceptions import Forbidden, NotFound from configs import dify_config from controllers.common.fields import SimpleResultResponse from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from controllers.common.wraps import enforce_rbac_access from controllers.console import console_ns +from controllers.console.flask_admission import console_account_admission from controllers.console.wraps import ( RBACPermission, RBACResourceScope, - account_initialization_required, - edit_permission_required, model_validate, - setup_required, - with_current_tenant_id, - with_current_user, ) -from extensions.ext_database import db +from extensions.ext_application_services import application_services from fields.base import ResponseModel from libs.helper import dump_response -from libs.login import current_account_with_tenant, login_required -from models import Account +from libs.login import current_account_with_tenant +from machinery.context import RequestContext from models.enums import TagType -from models.model import Tag -from services.tag_service import ( - SaveTagPayload, - TagBindingCreatePayload, - TagBindingDeletePayload, - TagService, - UpdateTagPayload, +from services.tag_application_service import ( + CreateTagInput, + TagBindingInput, + TagBindingTargetNotFoundError, + TagNameConflictError, + TagNotFoundError, + UpdateTagInput, ) @@ -59,7 +54,7 @@ class TagBindingRemovePayload(BaseModel): class TagListQueryParam(BaseModel): - type: Literal["knowledge", "app", "snippet", ""] = Field("", description="Tag type filter") + type: Literal["knowledge", "app", "snippet"] = Field(description="Tag type filter") keyword: str | None = Field(None, description="Search keyword") @@ -101,143 +96,158 @@ register_schema_models( register_response_schema_models(console_ns, SimpleResultResponse, TagResponse, TagListResponse) -def _enforce_snippet_tag_rbac_if_needed(tag_type: TagType | str | None) -> None: +def _enforce_snippet_tag_rbac_if_needed(tag_type: TagType | str | None, context: RequestContext) -> None: if tag_type != TagType.SNIPPET: return if not dify_config.RBAC_ENABLED: return - current_user, current_tenant_id = current_account_with_tenant() enforce_rbac_access( - tenant_id=current_tenant_id, - account_id=current_user.id, + tenant_id=_workspace_id(context), + account_id=context.account_id, resource_type=RBACResourceScope.WORKSPACE, scene=RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False, ) -def _enforce_snippet_tag_rbac_by_tag_id(tag_id: str) -> None: +def _enforce_snippet_tag_rbac_by_tag_id(tag_id: str, context: RequestContext) -> None: if not dify_config.RBAC_ENABLED: return - _, current_tenant_id = current_account_with_tenant() - tag_type = db.session.scalar(select(Tag.type).where(Tag.id == tag_id, Tag.tenant_id == current_tenant_id).limit(1)) - _enforce_snippet_tag_rbac_if_needed(tag_type) + tag_type = application_services().tags.get_tag_type(context, tag_id) + _enforce_snippet_tag_rbac_if_needed(tag_type, context) + + +def _workspace_id(context: RequestContext) -> str: + if context.active_workspace_id is None: + raise RuntimeError("Console account admission did not resolve an active workspace") + return context.active_workspace_id + + +def _require_tag_edit_permission(*, allow_dataset_editor: bool) -> None: + current_user, _ = current_account_with_tenant() + if current_user.has_edit_permission: + return + if allow_dataset_editor and current_user.is_dataset_editor: + return + raise Forbidden() @console_ns.route("/tags") class TagListApi(Resource): - @setup_required - @login_required - @account_initialization_required + @console_account_admission() @console_ns.doc(params=query_params_from_model(TagListQueryParam)) @console_ns.response(200, "Success", console_ns.models[TagListResponse.__name__]) - @with_current_tenant_id @model_validate(TagListQueryParam) - def get(self, req_data: TagListQueryParam, current_tenant_id: str): - tags = TagService.get_tags(req_data.type, current_tenant_id, req_data.keyword, session=db.session()) + def get(self, req_data: TagListQueryParam, request_context: RequestContext): + tags = application_services().tags.list_tags(request_context, req_data.type, req_data.keyword) return dump_response(TagListResponse, tags), 200 @console_ns.expect(console_ns.models[TagBasePayload.__name__]) @console_ns.response(200, "Success", console_ns.models[TagResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user + @console_account_admission() @model_validate(TagBasePayload) - def post(self, req_data: TagBasePayload, current_user: Account): + def post(self, req_data: TagBasePayload, request_context: RequestContext): # Allow users with edit permission, or dataset editors (including dataset operators). - if not (current_user.has_edit_permission or current_user.is_dataset_editor): - raise Forbidden() + _require_tag_edit_permission(allow_dataset_editor=True) - _enforce_snippet_tag_rbac_if_needed(req_data.type) - tag = TagService.save_tags(SaveTagPayload(name=req_data.name, type=req_data.type), db.session()) + _enforce_snippet_tag_rbac_if_needed(req_data.type, request_context) + try: + tag = application_services().tags.create_tag( + request_context, + CreateTagInput(name=req_data.name, type=req_data.type.value), + ) + except TagNameConflictError as error: + raise ValueError(str(error)) from None - return dump_response(TagResponse, {"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": 0}), 200 + return dump_response(TagResponse, tag), 200 @console_ns.route("/tags/") class TagUpdateDeleteApi(Resource): @console_ns.expect(console_ns.models[TagUpdateRequestPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[TagResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user + @console_account_admission() @model_validate(TagUpdateRequestPayload) - def patch(self, req_data: TagUpdateRequestPayload, current_user: Account, tag_id: UUID): + def patch(self, req_data: TagUpdateRequestPayload, request_context: RequestContext, tag_id: UUID): tag_id_str = str(tag_id) # The role of the current user in the ta table must be admin, owner, or editor - if not (current_user.has_edit_permission or current_user.is_dataset_editor): - raise Forbidden() + _require_tag_edit_permission(allow_dataset_editor=True) - _enforce_snippet_tag_rbac_by_tag_id(tag_id_str) - tag = TagService.update_tags(UpdateTagPayload(name=req_data.name), tag_id_str, db.session()) + _enforce_snippet_tag_rbac_by_tag_id(tag_id_str, request_context) + try: + tag = application_services().tags.update_tag( + request_context, + tag_id_str, + UpdateTagInput(name=req_data.name), + ) + except TagNameConflictError as error: + raise ValueError(str(error)) from None + except TagNotFoundError as error: + raise NotFound(str(error)) from None - binding_count = TagService.get_tag_binding_count(tag_id_str, db.session()) + return dump_response(TagResponse, tag), 200 - return ( - dump_response( - TagResponse, - {"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": binding_count}, - ), - 200, - ) - - @setup_required - @login_required - @account_initialization_required - @edit_permission_required + @console_account_admission() @console_ns.response(204, "Tag deleted successfully") - def delete(self, tag_id: UUID): + def delete(self, request_context: RequestContext, tag_id: UUID): tag_id_str = str(tag_id) - _enforce_snippet_tag_rbac_by_tag_id(tag_id_str) - TagService.delete_tag(tag_id_str, db.session()) + _require_tag_edit_permission(allow_dataset_editor=False) + _enforce_snippet_tag_rbac_by_tag_id(tag_id_str, request_context) + try: + application_services().tags.delete_tag(request_context, tag_id_str) + except TagNotFoundError as error: + raise NotFound(str(error)) from None return "", 204 -def _require_tag_binding_edit_permission(current_user: Account) -> None: +def _require_tag_binding_edit_permission() -> None: """ Ensure the current account can edit tag bindings. Tag binding operations are allowed for users who can edit resources (app/dataset) within the current tenant. """ # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator - if not (current_user.has_edit_permission or current_user.is_dataset_editor): - raise Forbidden() + _require_tag_edit_permission(allow_dataset_editor=True) -def _create_tag_bindings(current_user: Account, payload: TagBindingPayload) -> tuple[dict[str, str], int]: - _require_tag_binding_edit_permission(current_user) +def _create_tag_bindings(context: RequestContext, payload: TagBindingPayload) -> tuple[dict[str, str], int]: + _require_tag_binding_edit_permission() - _enforce_snippet_tag_rbac_if_needed(payload.type) - TagService.save_tag_binding( - TagBindingCreatePayload( - tag_ids=payload.tag_ids, - target_id=payload.target_id, - type=payload.type, - ), - db.session(), - ) + _enforce_snippet_tag_rbac_if_needed(payload.type, context) + try: + application_services().tags.create_bindings( + context, + TagBindingInput( + tag_ids=tuple(payload.tag_ids), + target_id=payload.target_id, + type=payload.type.value, + ), + ) + except TagBindingTargetNotFoundError as error: + raise NotFound(str(error)) from None return {"result": "success"}, 200 -def _remove_tag_bindings(current_user: Account, payload: TagBindingRemovePayload) -> tuple[dict[str, str], int]: - _require_tag_binding_edit_permission(current_user) +def _remove_tag_bindings(context: RequestContext, payload: TagBindingRemovePayload) -> tuple[dict[str, str], int]: + _require_tag_binding_edit_permission() - _enforce_snippet_tag_rbac_if_needed(payload.type) - TagService.delete_tag_binding( - TagBindingDeletePayload( - tag_ids=payload.tag_ids, - target_id=payload.target_id, - type=payload.type, - ), - db.session(), - ) + _enforce_snippet_tag_rbac_if_needed(payload.type, context) + try: + application_services().tags.delete_bindings( + context, + TagBindingInput( + tag_ids=tuple(payload.tag_ids), + target_id=payload.target_id, + type=payload.type.value, + ), + ) + except TagBindingTargetNotFoundError as error: + raise NotFound(str(error)) from None return {"result": "success"}, 200 @@ -248,13 +258,10 @@ class TagBindingCollectionApi(Resource): @console_ns.doc("create_tag_binding") @console_ns.expect(console_ns.models[TagBindingPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user + @console_account_admission() @model_validate(TagBindingPayload) - def post(self, req_data: TagBindingPayload, current_user: Account): - return _create_tag_bindings(current_user, req_data) + def post(self, req_data: TagBindingPayload, request_context: RequestContext): + return _create_tag_bindings(request_context, req_data) @console_ns.route("/tag-bindings/remove") @@ -265,10 +272,7 @@ class TagBindingRemoveApi(Resource): @console_ns.doc(description="Remove one or more tag bindings from a target.") @console_ns.expect(console_ns.models[TagBindingRemovePayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user + @console_account_admission() @model_validate(TagBindingRemovePayload) - def post(self, req_data: TagBindingRemovePayload, current_user: Account): - return _remove_tag_bindings(current_user, req_data) + def post(self, req_data: TagBindingRemovePayload, request_context: RequestContext): + return _remove_tag_bindings(request_context, req_data) diff --git a/api/controllers/console/workspace/account.py b/api/controllers/console/workspace/account.py index 9ef75cc03f8..5ce0494be80 100644 --- a/api/controllers/console/workspace/account.py +++ b/api/controllers/console/workspace/account.py @@ -2,12 +2,13 @@ from __future__ import annotations from datetime import datetime from http import HTTPStatus -from typing import Literal +from typing import Annotated, Literal import pytz from flask import request from flask_restx import Resource -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic.json_schema import SkipJsonSchema from sqlalchemy import select from werkzeug.exceptions import NotFound @@ -28,7 +29,13 @@ from controllers.console.auth.error import ( InvalidEmailError, InvalidTokenError, ) -from controllers.console.error import AccountInFreezeError, AccountNotFound, EmailSendIpLimitError +from controllers.console.error import ( + AccountInFreezeError, + AccountNotFound, + EmailDomainSuspendedError, + EmailSendIpLimitError, +) +from controllers.console.flask_admission import console_account_admission from controllers.console.workspace.error import ( AccountAlreadyInitedError, CurrentPasswordIncorrectError, @@ -46,6 +53,7 @@ from controllers.console.wraps import ( with_current_user, ) from enums import DeploymentEdition +from extensions.ext_application_services import application_services from extensions.ext_database import db from fields.base import ResponseModel from fields.member_fields import AccountResponse @@ -53,12 +61,15 @@ from graphon.file import helpers as file_helpers from libs.datetime_utils import naive_utc_now from libs.helper import EmailStr, dump_response, extract_remote_ip, timezone, to_timestamp from libs.login import login_required +from machinery.context import RequestContext from models import Account, AccountIntegrate, InvitationCode from models.account import AccountStatus, InvitationCodeStatus from models.enums import CreatorUserRole from models.model import UploadFile +from services import account_errors from services.account_service import AccountService from services.billing_service import BillingService +from services.entities.account_entities import AccountProfileChanges from services.entities.auth_entities import ( ChangeEmailNewEmailToken, ChangeEmailNewEmailVerifiedToken, @@ -118,6 +129,42 @@ class AccountTimezonePayload(BaseModel): return timezone(value) +class AccountProfilePatchPayload(BaseModel): + model_config = ConfigDict(extra="forbid") + + name: Annotated[str, Field(min_length=3, max_length=30)] | SkipJsonSchema[None] = None + avatar: str | SkipJsonSchema[None] = None + interface_language: str | SkipJsonSchema[None] = None + interface_theme: Literal["light", "dark"] | SkipJsonSchema[None] = None + timezone: str | SkipJsonSchema[None] = None + + @field_validator("*", mode="before") + @classmethod + def reject_null(cls, value: object) -> object: + if value is None: + raise ValueError("Account profile fields cannot be null") + return value + + @field_validator("interface_language") + @classmethod + def validate_language(cls, value: str) -> str: + return supported_language(value) + + @field_validator("timezone") + @classmethod + def validate_timezone(cls, value: str) -> str: + return timezone(value) + + def to_changes(self) -> AccountProfileChanges: + return AccountProfileChanges( + name=self.name, + avatar=self.avatar, + interface_language=self.interface_language, + interface_theme=self.interface_theme, + timezone=self.timezone, + ) + + class AccountPasswordPayload(BaseModel): password: str | None = None new_password: str @@ -183,6 +230,7 @@ register_schema_models( AccountInterfaceLanguagePayload, AccountInterfaceThemePayload, AccountTimezonePayload, + AccountProfilePatchPayload, AccountPasswordPayload, AccountDeletePayload, AccountDeletionFeedbackPayload, @@ -248,6 +296,14 @@ register_response_schema_models( ) +def _update_account_profile(request_context: RequestContext, changes: AccountProfileChanges) -> dict[str, object]: + try: + account = application_services().accounts.profile.update(request_context, changes) + except account_errors.AccountNotFoundError as error: + raise AccountNotFound() from error + return dump_response(AccountResponse, account) + + @console_ns.route("/account/init") class AccountInitApi(Resource): @console_ns.expect(console_ns.models[AccountInitPayload.__name__]) @@ -305,21 +361,28 @@ class AccountProfileApi(Resource): def get(self, current_user: Account): return dump_response(AccountResponse, current_user) + @console_ns.expect(console_ns.models[AccountProfilePatchPayload.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) + @console_account_admission() + @model_validate(AccountProfilePatchPayload) + def patch(self, args: AccountProfilePatchPayload, request_context: RequestContext): + return _update_account_profile(request_context, args.to_changes()) + @console_ns.route("/account/name") class AccountNameApi(Resource): + """Deprecated compatibility route; use PATCH /account/profile.""" + + @console_ns.doc("update_account_name_deprecated") + @console_ns.doc(deprecated=True) + @console_ns.doc(description="Deprecated. Use PATCH /account/profile instead.") @console_ns.expect(console_ns.models[AccountNamePayload.__name__]) - @setup_required - @login_required - @account_initialization_required @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) - @with_current_user - def post(self, current_user: Account): + @console_account_admission() + def post(self, request_context: RequestContext): payload = console_ns.payload or {} args = AccountNamePayload.model_validate(payload) - updated_account = AccountService.update_account(current_user, session=db.session(), name=args.name) - - return dump_response(AccountResponse, updated_account) + return _update_account_profile(request_context, AccountProfileChanges(name=args.name)) @console_ns.route("/account/avatar") @@ -350,73 +413,69 @@ class AccountAvatarApi(Resource): return AvatarUrlResponse(avatar_url=avatar_url).model_dump(mode="json") @console_ns.expect(console_ns.models[AccountAvatarPayload.__name__]) - @setup_required - @login_required - @account_initialization_required + @console_ns.doc("update_account_avatar_deprecated") + @console_ns.doc(deprecated=True) + @console_ns.doc(description="Deprecated. Use PATCH /account/profile instead.") @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) - @with_current_user - def post(self, current_user: Account): + @console_account_admission() + def post(self, request_context: RequestContext): payload = console_ns.payload or {} args = AccountAvatarPayload.model_validate(payload) - - updated_account = AccountService.update_account(current_user, session=db.session(), avatar=args.avatar) - - return dump_response(AccountResponse, updated_account) + return _update_account_profile(request_context, AccountProfileChanges(avatar=args.avatar)) @console_ns.route("/account/interface-language") class AccountInterfaceLanguageApi(Resource): + """Deprecated compatibility route; use PATCH /account/profile.""" + + @console_ns.doc("update_account_interface_language_deprecated") + @console_ns.doc(deprecated=True) + @console_ns.doc(description="Deprecated. Use PATCH /account/profile instead.") @console_ns.expect(console_ns.models[AccountInterfaceLanguagePayload.__name__]) - @setup_required - @login_required - @account_initialization_required @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) - @with_current_user - def post(self, current_user: Account): + @console_account_admission() + def post(self, request_context: RequestContext): payload = console_ns.payload or {} args = AccountInterfaceLanguagePayload.model_validate(payload) - - updated_account = AccountService.update_account( - current_user, session=db.session(), interface_language=args.interface_language + return _update_account_profile( + request_context, + AccountProfileChanges(interface_language=args.interface_language), ) - return dump_response(AccountResponse, updated_account) - @console_ns.route("/account/interface-theme") class AccountInterfaceThemeApi(Resource): + """Deprecated compatibility route; use PATCH /account/profile.""" + + @console_ns.doc("update_account_interface_theme_deprecated") + @console_ns.doc(deprecated=True) + @console_ns.doc(description="Deprecated. Use PATCH /account/profile instead.") @console_ns.expect(console_ns.models[AccountInterfaceThemePayload.__name__]) - @setup_required - @login_required - @account_initialization_required @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) - @with_current_user - def post(self, current_user: Account): + @console_account_admission() + def post(self, request_context: RequestContext): payload = console_ns.payload or {} args = AccountInterfaceThemePayload.model_validate(payload) - - updated_account = AccountService.update_account( - current_user, session=db.session(), interface_theme=args.interface_theme + return _update_account_profile( + request_context, + AccountProfileChanges(interface_theme=args.interface_theme), ) - return dump_response(AccountResponse, updated_account) - @console_ns.route("/account/timezone") class AccountTimezoneApi(Resource): + """Deprecated compatibility route; use PATCH /account/profile.""" + + @console_ns.doc("update_account_timezone_deprecated") + @console_ns.doc(deprecated=True) + @console_ns.doc(description="Deprecated. Use PATCH /account/profile instead.") @console_ns.expect(console_ns.models[AccountTimezonePayload.__name__]) - @setup_required - @login_required - @account_initialization_required @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) - @with_current_user - def post(self, current_user: Account): + @console_account_admission() + def post(self, request_context: RequestContext): payload = console_ns.payload or {} args = AccountTimezonePayload.model_validate(payload) - - updated_account = AccountService.update_account(current_user, session=db.session(), timezone=args.timezone) - - return dump_response(AccountResponse, updated_account) + return _update_account_profile(request_context, AccountProfileChanges(timezone=args.timezone)) @console_ns.route("/account/password") @@ -715,7 +774,10 @@ class ChangeEmailResetApi(Resource): args = ChangeEmailResetPayload.model_validate(payload) normalized_new_email = args.new_email.lower() - if AccountService.is_account_in_freeze(normalized_new_email): + freeze_type = AccountService.get_account_freeze_type(normalized_new_email) + if freeze_type: + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedError() raise AccountInFreezeError() if not AccountService.check_email_unique(normalized_new_email, session=db.session()): @@ -762,7 +824,10 @@ class CheckEmailUnique(Resource): payload = console_ns.payload or {} args = CheckEmailUniquePayload.model_validate(payload) normalized_email = args.email.lower() - if AccountService.is_account_in_freeze(normalized_email): + freeze_type = AccountService.get_account_freeze_type(normalized_email) + if freeze_type: + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedError() raise AccountInFreezeError() if not AccountService.check_email_unique(normalized_email, session=db.session()): raise EmailAlreadyInUseError() diff --git a/api/controllers/console/workspace/endpoint.py b/api/controllers/console/workspace/endpoint.py index f02ae4a73fb..ef87c6c4571 100644 --- a/api/controllers/console/workspace/endpoint.py +++ b/api/controllers/console/workspace/endpoint.py @@ -290,6 +290,8 @@ class EndpointListApi(Resource): ) @setup_required @login_required + @is_admin_or_owner_required + @rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_MODEL_CONFIG, resource_required=False) @account_initialization_required @with_current_user_id @with_current_tenant_id @@ -318,6 +320,8 @@ class EndpointListForSinglePluginApi(Resource): ) @setup_required @login_required + @is_admin_or_owner_required + @rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_MODEL_CONFIG, resource_required=False) @account_initialization_required @with_current_user_id @with_current_tenant_id diff --git a/api/controllers/console/workspace/snippets.py b/api/controllers/console/workspace/snippets.py index 8c254dfa58b..f5e2f85a0b7 100644 --- a/api/controllers/console/workspace/snippets.py +++ b/api/controllers/console/workspace/snippets.py @@ -17,7 +17,7 @@ from controllers.console import console_ns from controllers.console.app.wraps import with_session from controllers.console.snippets.payloads import ( CreateSnippetPayload, - IncludeSecretQuery, + SnippetExportQuery, SnippetImportPayload, SnippetListQuery, UpdateSnippetPayload, @@ -89,7 +89,7 @@ register_schema_models( CreateSnippetPayload, UpdateSnippetPayload, SnippetImportPayload, - IncludeSecretQuery, + SnippetExportQuery, ) register_response_schema_models( console_ns, @@ -289,7 +289,7 @@ class CustomizedSnippetExportApi(Resource): @console_ns.doc("export_customized_snippet") @console_ns.doc(description="Export snippet configuration as DSL") @console_ns.doc(params={"snippet_id": "Snippet ID to export"}) - @console_ns.doc(params=query_params_from_model(IncludeSecretQuery)) + @console_ns.doc(params=query_params_from_model(SnippetExportQuery)) @console_ns.response(200, "Snippet exported successfully", console_ns.models[TextFileResponse.__name__]) @console_ns.response(404, "Snippet not found") @setup_required @@ -312,11 +312,18 @@ class CustomizedSnippetExportApi(Resource): raise NotFound("Snippet not found") # Get include_secret parameter - query = IncludeSecretQuery.model_validate(request.args.to_dict()) + query = SnippetExportQuery.model_validate(request.args.to_dict()) with Session(db.engine) as session: export_service = SnippetDslService(session) - result = export_service.export_snippet_dsl(snippet=snippet, include_secret=query.include_secret == "true") + try: + result = export_service.export_snippet_dsl( + snippet=snippet, + include_secret=query.include_secret == "true", + workflow_id=query.workflow_id, + ) + except ValueError as exc: + raise NotFound(str(exc)) from exc # Set filename with .snippet extension filename = f"{snippet.name}.snippet" diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py index a02e603373d..cb61f1582b6 100644 --- a/api/controllers/console/workspace/tool_providers.py +++ b/api/controllers/console/workspace/tool_providers.py @@ -1296,6 +1296,8 @@ class ToolOAuthCustomClient(Resource): ) @setup_required @login_required + @is_admin_or_owner_required + @rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.CREDENTIAL_MANAGE, resource_required=False) @account_initialization_required @with_current_tenant_id def get(self, current_tenant_id: str, provider: str): diff --git a/api/controllers/openapi/app_dsl.py b/api/controllers/openapi/app_dsl.py index d06845dada4..5d036580654 100644 --- a/api/controllers/openapi/app_dsl.py +++ b/api/controllers/openapi/app_dsl.py @@ -4,6 +4,7 @@ from typing import cast from flask_restx import Resource from sqlalchemy.orm import Session +from werkzeug.exceptions import Forbidden from controllers.common.wraps import RBACPermission, RBACResourceScope from controllers.openapi import openapi_ns @@ -17,6 +18,7 @@ from models import Account, App from models.account import TenantAccountRole from services.app_dsl_service import AppDslService, Import from services.entities.dsl_entities import CheckDependenciesResult, ImportStatus +from services.errors.account import NoPermissionError from services.errors.app import WorkflowNotFoundError @@ -53,18 +55,21 @@ class AppDslImportApi(Resource): with Session(db.engine, expire_on_commit=False) as session: service = AppDslService(session) - result = service.import_app( - account=account, - import_mode=body.mode, - yaml_content=body.yaml_content, - yaml_url=body.yaml_url, - name=body.name, - description=body.description, - icon_type=body.icon_type, - icon=body.icon, - icon_background=body.icon_background, - app_id=body.app_id, - ) + try: + result = service.import_app( + account=account, + import_mode=body.mode, + yaml_content=body.yaml_content, + yaml_url=body.yaml_url, + name=body.name, + description=body.description, + icon_type=body.icon_type, + icon=body.icon, + icon_background=body.icon_background, + app_id=body.app_id, + ) + except NoPermissionError as exc: + raise Forbidden(str(exc)) from exc if result.status == ImportStatus.FAILED: session.rollback() else: @@ -108,7 +113,10 @@ class AppDslImportConfirmApi(Resource): with Session(db.engine, expire_on_commit=False) as session: service = AppDslService(session) - result = service.confirm_import(import_id=import_id, account=account) + try: + result = service.confirm_import(import_id=import_id, account=account) + except NoPermissionError as exc: + raise Forbidden(str(exc)) from exc if result.status == ImportStatus.FAILED: session.rollback() else: diff --git a/api/controllers/web/app.py b/api/controllers/web/app.py index 135765a902c..1349de9d486 100644 --- a/api/controllers/web/app.py +++ b/api/controllers/web/app.py @@ -7,33 +7,29 @@ from pydantic import BaseModel, ConfigDict, Field from werkzeug.exceptions import Unauthorized from constants import HEADER_NAME_APP_CODE -from controllers.common import fields from controllers.common.errors import InvalidArgumentError -from controllers.common.fields import AccessModeResponse, Parameters +from controllers.common.fields import AccessModeResponse, BooleanResultResponse, Parameters from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from controllers.web import web_ns from controllers.web.error import ( AgentNotPublishedError, AppUnavailableError, WebAppAccessServiceUnavailableError, + WebAppAuthRequiredError, WebAppNotFoundError, ) from controllers.web.wraps import WebApiResource from extensions.ext_application_services import application_services -from extensions.ext_database import db from libs.helper import dump_response from libs.passport import PassportService from libs.token import extract_webapp_passport from models.model import App, EndUser from services.app_definition_query_service import AppDefinitionNotPublishedError, AppDefinitionUnavailableError -from services.enterprise.enterprise_service import EnterpriseService -from services.feature_service import FeatureService from services.webapp_access_query_service import ( WebAppAccessAppNotFoundError, WebAppAccessReferenceRequiredError, WebAppAccessUnavailableError, ) -from services.webapp_auth_service import WebAppAuthService logger = logging.getLogger(__name__) @@ -64,7 +60,7 @@ register_response_schema_models( Parameters, AppMetaResponse, AccessModeResponse, - fields.BooleanResultResponse, + BooleanResultResponse, ) @@ -165,21 +161,23 @@ class AppWebAuthPermission(Resource): 400: "Bad Request", 401: "Unauthorized", 500: "Internal Server Error", + 503: "Web App Access Service Unavailable", } ) - @web_ns.response(200, "Success", web_ns.models[fields.BooleanResultResponse.__name__]) + @web_ns.response(200, "Success", web_ns.models[BooleanResultResponse.__name__]) def get(self): - user_id = "visitor" app_code = request.headers.get(HEADER_NAME_APP_CODE) app_id = request.args.get("appId") if not app_id or not app_code: raise ValueError("appId must be provided") - require_permission_check = WebAppAuthService.is_app_require_permission_check( - app_id=app_id, session=db.session() - ) - if not require_permission_check: - return {"result": True} + webapp_access = application_services().webapp_access + try: + requires_permission_check = webapp_access.requires_permission_check(app_id) + except WebAppAccessUnavailableError: + raise WebAppAccessServiceUnavailableError() from None + if not requires_permission_check: + return dump_response(BooleanResultResponse, {"result": True}) try: tk = extract_webapp_passport(app_code, request) @@ -188,16 +186,13 @@ class AppWebAuthPermission(Resource): decoded = PassportService().verify(tk) user_id = decoded.get("user_id", "visitor") except Unauthorized: - raise + raise WebAppAuthRequiredError() from None except Exception: logger.exception("Unexpected error during auth verification") raise - features = FeatureService.get_system_features() - if not features.webapp_auth.enabled: - return {"result": True} - - res = True - if WebAppAuthService.is_app_require_permission_check(app_id=app_id, session=db.session()): - res = EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp(str(user_id), app_id) - return {"result": res} + try: + is_allowed = webapp_access.is_user_allowed(user_id=str(user_id), app_id=app_id) + except WebAppAccessUnavailableError: + raise WebAppAccessServiceUnavailableError() from None + return dump_response(BooleanResultResponse, {"result": is_allowed}) diff --git a/api/controllers/web/site.py b/api/controllers/web/site.py index 827fdce72b0..54914b9bfb4 100644 --- a/api/controllers/web/site.py +++ b/api/controllers/web/site.py @@ -1,23 +1,19 @@ from typing import Any, Self from pydantic import AliasChoices, Field -from sqlalchemy import select from werkzeug.exceptions import Forbidden from configs import dify_config from controllers.common.schema import register_response_schema_models from controllers.web import web_ns from controllers.web.wraps import WebApiResource -from enums import DeploymentEdition -from extensions.ext_database import db -from extensions.storage.storage_type import StorageType +from extensions.ext_application_services import application_services from fields.base import ResponseModel -from libs.helper import build_icon_url -from models.account import Tenant, TenantStatus -from models.model import App, AppMode, EndUser, IconType, Site +from libs.helper import build_icon_url, dump_response +from models.account import Tenant +from models.model import App, AppMode, EndUser, Site from services.entities.feature_entities import FeatureModel -from services.feature_service import FeatureService -from services.file_service import FileService +from services.web_app_runtime_query_service import WebAppRuntimeUnavailableError class WebSiteResponse(ResponseModel): @@ -128,17 +124,6 @@ register_response_schema_models( ) -def _build_site_icon_url(*, site: Site, tenant_id: str) -> str | None: - """Use direct S3 URLs only in Cloud Mode and preserve preview URLs elsewhere.""" - if site.icon_type != IconType.IMAGE or not site.icon: - return None - if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and ( - StorageType(dify_config.STORAGE_TYPE) == StorageType.S3 - ): - return FileService(db.engine).get_file_presigned_url(file_id=site.icon, tenant_id=tenant_id) - return build_icon_url(site.icon_type, site.icon) - - @web_ns.route("/site") class AppSiteApi(WebApiResource): @web_ns.doc("Get App Site Info") @@ -156,25 +141,21 @@ class AppSiteApi(WebApiResource): @web_ns.response(200, "Success", web_ns.models[WebAppSiteResponse.__name__]) def get(self, app_model: App, end_user: EndUser): """Retrieve app site info.""" - # get site - site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1)) + try: + bootstrap = application_services().web_app_runtime.get_bootstrap(app_model.id) + except WebAppRuntimeUnavailableError: + raise Forbidden() from None - if site is None: - raise Forbidden() - - tenant = app_model.tenant - if tenant is None or tenant.status == TenantStatus.ARCHIVE: - raise Forbidden() - - features = FeatureService.get_features(app_model.tenant_id, exclude_vector_space=True) - - return WebAppSiteResponse.from_app_site( - tenant=tenant, - app_model=app_model, - mode=AppMode.value_of(app_model.mode_compatible_with_agent_with_session(session=db.session())), - site=site, - end_user_id=end_user.id, - features=features, - can_replace_logo=features.can_replace_logo, - icon_url=_build_site_icon_url(site=site, tenant_id=tenant.id), - ).model_dump(mode="json") + return dump_response( + WebAppSiteResponse, + { + "app_id": bootstrap.app_id, + "mode": bootstrap.mode, + "end_user_id": end_user.id, + "enable_site": bootstrap.enable_site, + "site": bootstrap.site, + "plan": bootstrap.plan, + "can_replace_logo": bootstrap.can_replace_logo, + "custom_config": bootstrap.custom_config, + }, + ) diff --git a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py b/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py index 41aa002ae27..42b42a7f273 100644 --- a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py +++ b/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py @@ -273,6 +273,17 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat with session_factory.create_session() as session: err = self.handle_error(event=event, session=session, message_id=self._message_id) session.commit() + + if trace_manager: + trace_manager.add_trace_task( + TraceTask( + TraceTaskName.MESSAGE_TRACE, + conversation_id=self._conversation_id, + message_id=self._message_id, + trace_session_id=self._application_generate_entity.extras.get("trace_session_id"), + ) + ) + yield self.error_to_stream_response(err) break case QueueStopEvent() | QueueMessageEndEvent(): diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index b91bf554897..37e4fcdcda9 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -16,10 +16,15 @@ from core.schemas.schema_manager import SchemaManager from enums import DeploymentEdition, WebAppAccessMode from extensions.ext_redis import RedisClientWrapper, redis_client from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository +from repositories.account_repository import SQLAlchemyAccountRepository from repositories.app_definition_query_repository import AppDefinitionQueryRepository from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository from repositories.explore_banner_query_repository import ExploreBannerQueryRepository from repositories.installation_state_repository import InstallationStateRepository +from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository +from repositories.tag_repository import TagRepository +from repositories.trial_app_query_repository import TrialAppQueryRepository +from repositories.trial_app_usage_repository import TrialAppUsageRepository from repositories.webapp_access_query_repository import WebAppAccessQueryRepository from repositories.workspace_member_query_repository import WorkspaceMemberQueryRepository from repositories.workspace_query_repository import WorkspaceQueryRepository @@ -30,6 +35,7 @@ from services.account_activation_adapters import ( RegisterServiceInvitationTokenStore, ) from services.account_activation_service import AccountActivationService +from services.account_profile_service import AccountProfileService from services.app_definition_query_service import AppDefinitionQueryService from services.auth.data_source_api_key_auth_gateways import ( ProviderApiKeyAuthCredentialValidator, @@ -42,10 +48,20 @@ from services.explore_banner_query_service import ExploreBannerQueryService from services.feature_query_service import FeatureQueryService from services.feature_service import FeatureService from services.feature_service_gateway import FeatureServiceGateway +from services.file_service import FileService from services.init_validation_service import InitValidationService +from services.recommended_app_catalog_gateway import ( + BuiltinRecommendedAppCatalogGateway, + RecommendedAppCatalogRouter, + RemoteRecommendedAppCatalogGateway, +) +from services.recommended_app_query_service import RecommendedAppQueryService from services.schema_definition_service import SchemaDefinitionService from services.setup_adapters import RedisSetupLock, RegisterServiceAccountProvisioner from services.setup_service import SetupService +from services.tag_application_service import TagApplicationService +from services.trial_app_usage import TrialAppUsageRecorder +from services.web_app_runtime_query_service import WebAppRuntimeQueryService from services.webapp_access_query_service import ( WebAppAccessQueryService, WebAppAccessUnavailableError, @@ -69,19 +85,36 @@ def _get_enterprise_webapp_access_mode(app_id: str) -> WebAppAccessMode: raise WebAppAccessUnavailableError from e +def _is_user_allowed_to_access_webapp(user_id: str, app_id: str) -> bool: + try: + return EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp(user_id, app_id) + except (EnterpriseServiceError, httpx.RequestError, json.JSONDecodeError, UnicodeDecodeError) as e: + raise WebAppAccessUnavailableError from e + + +@dataclass(frozen=True, slots=True) +class AccountServices: + profile: AccountProfileService + + @dataclass(frozen=True, slots=True) class ApplicationServices: + accounts: AccountServices account_activation: AccountActivationService app_definitions: AppDefinitionQueryService data_source_api_key_auth: DataSourceApiKeyAuthService webapp_access: WebAppAccessQueryService + web_app_runtime: WebAppRuntimeQueryService explore_banner_queries: ExploreBannerQueryService schema_definitions: SchemaDefinitionService setup: SetupService feature_queries: FeatureQueryService init_validation: InitValidationService + recommended_app_queries: RecommendedAppQueryService + trial_app_usage: TrialAppUsageRecorder workspace_queries: WorkspaceQueryService workspace_member_queries: WorkspaceMemberQueryService + tags: TagApplicationService def build_application_services( @@ -93,7 +126,21 @@ def build_application_services( ) -> ApplicationServices: installation_state = InstallationStateRepository(client=database_client) data_source_api_key_auth_bindings = SQLAlchemyDataSourceApiKeyAuthBindingRepository(session_factory=database_client) + app_definition_repository = AppDefinitionQueryRepository(session_factory=database_client) + feature_gateway = FeatureServiceGateway() + trial_app_enabled = FeatureService.is_trial_app_enabled() + database_catalog = DatabaseRecommendedAppCatalogRepository(session_factory=database_client, redis=redis) + builtin_catalog = BuiltinRecommendedAppCatalogGateway() + remote_catalog = RemoteRecommendedAppCatalogGateway() + recommended_app_catalog = RecommendedAppCatalogRouter( + remote=remote_catalog, + database=database_catalog, + builtin=builtin_catalog, + ) return ApplicationServices( + accounts=AccountServices( + profile=AccountProfileService(accounts=SQLAlchemyAccountRepository(database_client)), + ), account_activation=AccountActivationService( tokens=RegisterServiceInvitationTokenStore(), accounts=SQLAlchemyAccountActivationRepository(database_client), @@ -106,7 +153,7 @@ def build_application_services( ), ), app_definitions=AppDefinitionQueryService( - definitions=AppDefinitionQueryRepository(session_factory=database_client), + definitions=app_definition_repository, builtin_icon_url_prefix=( dify_config.CONSOLE_API_URL + "/console/api/workspaces/current/tool-provider/builtin/" ), @@ -120,6 +167,13 @@ def build_application_services( access=WebAppAccessQueryRepository(session_factory=database_client), webapp_auth_enabled=FeatureService.is_webapp_auth_enabled(), access_mode_for_app=_get_enterprise_webapp_access_mode, + is_user_allowed_for_app=_is_user_allowed_to_access_webapp, + ), + web_app_runtime=WebAppRuntimeQueryService( + runtime=app_definition_repository, + file_service=FileService(database_client), + workspace_features=feature_gateway.get_workspace_features, + files_url=dify_config.FILES_URL, ), explore_banner_queries=ExploreBannerQueryService( banners=ExploreBannerQueryRepository(client=database_client), @@ -133,7 +187,7 @@ def build_application_services( setup_required=deployment_edition != DeploymentEdition.CLOUD, ), feature_queries=FeatureQueryService( - features=FeatureServiceGateway(), + features=feature_gateway, trial_models=FeatureService.get_trial_models(), app_dsl_version=CURRENT_APP_DSL_VERSION, ), @@ -142,6 +196,12 @@ def build_application_services( validation_required=(deployment_edition != DeploymentEdition.CLOUD and bool(initialization_password)), expected_password=initialization_password, ), + recommended_app_queries=RecommendedAppQueryService( + catalog=recommended_app_catalog, + trial_apps=TrialAppQueryRepository(session_factory=database_client), + trial_enabled=trial_app_enabled, + ), + trial_app_usage=TrialAppUsageRepository(session_factory=database_client), workspace_queries=WorkspaceQueryService( workspaces=WorkspaceQueryRepository( client=database_client, @@ -154,6 +214,9 @@ def build_application_services( ), roles=DeploymentWorkspaceMemberRoleResolver(), ), + tags=TagApplicationService( + tags=TagRepository(session_factory=database_client), + ), ) diff --git a/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py b/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py index 6a2223697e3..7803297aa11 100644 --- a/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_api_workflow_node_execution_repository.py @@ -8,7 +8,7 @@ WorkflowNodeExecutionModel operations using Aliyun SLS LogStore. import logging import time from collections.abc import Mapping, Sequence -from datetime import datetime +from datetime import UTC, datetime from typing import Any, override from sqlalchemy.orm import sessionmaker @@ -17,6 +17,7 @@ from extensions.logstore.aliyun_logstore import AliyunLogStore from extensions.logstore.repositories import safe_float, safe_int from extensions.logstore.sql_escape import escape_identifier, escape_logstore_query_value from graphon.enums import WorkflowNodeExecutionStatus +from libs.datetime_utils import ensure_naive_utc, naive_utc_now from models.enums import CreatorUserRole from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository @@ -86,28 +87,29 @@ def _dict_to_workflow_node_execution_model(data: dict[str, Any]) -> WorkflowNode model.execution_metadata = data.get("execution_metadata") # Handle datetime fields + # Every branch must yield naive UTC, matching what the database path stores. created_at = data.get("created_at") match created_at: case None: # Provide default created_at if missing - model.created_at = datetime.now() + model.created_at = naive_utc_now() case str(): - model.created_at = datetime.fromisoformat(created_at) + model.created_at = ensure_naive_utc(datetime.fromisoformat(created_at)) case int() | float(): - model.created_at = datetime.fromtimestamp(created_at) + model.created_at = datetime.fromtimestamp(created_at, tz=UTC).replace(tzinfo=None) case _: - model.created_at = created_at + model.created_at = ensure_naive_utc(created_at) finished_at = data.get("finished_at") match finished_at: case None: ... case str(): - model.finished_at = datetime.fromisoformat(finished_at) + model.finished_at = ensure_naive_utc(datetime.fromisoformat(finished_at)) case int() | float(): - model.finished_at = datetime.fromtimestamp(finished_at) + model.finished_at = datetime.fromtimestamp(finished_at, tz=UTC).replace(tzinfo=None) case _: - model.finished_at = finished_at + model.finished_at = ensure_naive_utc(finished_at) return model diff --git a/api/extensions/logstore/repositories/logstore_api_workflow_run_repository.py b/api/extensions/logstore/repositories/logstore_api_workflow_run_repository.py index 66028fb85bb..40c638f9d35 100644 --- a/api/extensions/logstore/repositories/logstore_api_workflow_run_repository.py +++ b/api/extensions/logstore/repositories/logstore_api_workflow_run_repository.py @@ -17,7 +17,7 @@ import logging import os import time from collections.abc import Sequence -from datetime import datetime +from datetime import UTC, datetime from typing import Any, cast, override from sqlalchemy.orm import sessionmaker @@ -26,6 +26,7 @@ from extensions.logstore.aliyun_logstore import AliyunLogStore from extensions.logstore.repositories import safe_float, safe_int from extensions.logstore.sql_escape import escape_identifier, escape_logstore_query_value, escape_sql_string from graphon.enums import WorkflowExecutionStatus +from libs.datetime_utils import ensure_naive_utc, naive_utc_now from libs.infinite_scroll_pagination import InfiniteScrollPagination from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.workflow import WorkflowRun, WorkflowType @@ -104,28 +105,30 @@ def _dict_to_workflow_run(data: dict[str, Any]) -> WorkflowRun: model.error = data.get("error_message") or data.get("error") # Handle datetime fields + # Every branch must yield naive UTC, matching what the database path stores. + # Mixing naive local time and aware values here breaks the elapsed_time subtraction below. started_at = data.get("started_at") or data.get("created_at") if started_at: match started_at: case str(): - model.created_at = datetime.fromisoformat(started_at) + model.created_at = ensure_naive_utc(datetime.fromisoformat(started_at)) case int() | float(): - model.created_at = datetime.fromtimestamp(started_at) + model.created_at = datetime.fromtimestamp(started_at, tz=UTC).replace(tzinfo=None) case _: - model.created_at = started_at + model.created_at = ensure_naive_utc(started_at) else: # Provide default created_at if missing - model.created_at = datetime.now() + model.created_at = naive_utc_now() finished_at = data.get("finished_at") if finished_at: match finished_at: case str(): - model.finished_at = datetime.fromisoformat(finished_at) + model.finished_at = ensure_naive_utc(datetime.fromisoformat(finished_at)) case int() | float(): - model.finished_at = datetime.fromtimestamp(finished_at) + model.finished_at = datetime.fromtimestamp(finished_at, tz=UTC).replace(tzinfo=None) case _: - model.finished_at = finished_at + model.finished_at = ensure_naive_utc(finished_at) # Compute elapsed_time from started_at and finished_at # LogStore doesn't store elapsed_time, it's computed in WorkflowExecution domain entity diff --git a/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py b/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py index 640b9c4ff66..55d66aae680 100644 --- a/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_workflow_execution_repository.py @@ -1,12 +1,12 @@ import json import logging -import os import time from typing import override from sqlalchemy.engine import Engine from sqlalchemy.orm import sessionmaker +from configs import dify_config from core.repositories.factory import WorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository from extensions.logstore.aliyun_logstore import AliyunLogStore @@ -71,14 +71,12 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository): triggered_from=triggered_from, ) - # Control flag for dual-write (write to both LogStore and SQL database) - # Set to True to enable dual-write for safe migration, False to use LogStore only - self._enable_dual_write = os.environ.get("LOGSTORE_DUAL_WRITE_ENABLED", "false").lower() == "true" + self._enable_dual_write = dify_config.LOGSTORE_DUAL_WRITE_ENABLED # Control flag for whether to write the `graph` field to LogStore. # If LOGSTORE_ENABLE_PUT_GRAPH_FIELD is "true", write the full `graph` field; # otherwise write an empty {} instead. Defaults to writing the `graph` field. - self._enable_put_graph_field = os.environ.get("LOGSTORE_ENABLE_PUT_GRAPH_FIELD", "true").lower() == "true" + self._enable_put_graph_field = dify_config.LOGSTORE_ENABLE_PUT_GRAPH_FIELD def _to_logstore_model(self, domain_model: WorkflowExecution) -> list[tuple[str, str]]: """ diff --git a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py index fa11e210c05..c75177750b3 100644 --- a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py @@ -7,7 +7,6 @@ using Aliyun SLS LogStore with append-only writes and version control. import json import logging -import os import time from collections.abc import Sequence from datetime import datetime @@ -16,6 +15,7 @@ from typing import Any, override from sqlalchemy.engine import Engine from sqlalchemy.orm import sessionmaker +from configs import dify_config from core.ops.utils import JSON_DICT_ADAPTER from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository from core.repositories.factory import OrderConfig, WorkflowNodeExecutionRepository @@ -152,9 +152,9 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): triggered_from=triggered_from, ) - # Control flag for dual-write (write to both LogStore and SQL database) - # Set to True to enable dual-write for safe migration, False to use LogStore only - self._enable_dual_write = os.environ.get("LOGSTORE_DUAL_WRITE_ENABLED", "false").lower() == "true" + # Keep the migration switch on the typed application config so callers + # and tests share the same validated source. + self._enable_dual_write = dify_config.LOGSTORE_DUAL_WRITE_ENABLED def _to_logstore_model(self, domain_model: WorkflowNodeExecution) -> Sequence[tuple[str, str]]: logger.debug( diff --git a/api/migrations/versions/2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py b/api/migrations/versions/2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py new file mode 100644 index 00000000000..32120abb0b2 --- /dev/null +++ b/api/migrations/versions/2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py @@ -0,0 +1,55 @@ +"""clean legacy agent soul files + +Revision ID: fbdfcf5f5a6e +Revises: 89919253ca7a +Create Date: 2026-08-20 09:38:36.827807 + +""" + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "fbdfcf5f5a6e" +down_revision = "89919253ca7a" +branch_labels = None +depends_on = None + + +def _remove_legacy_files(table_name: str) -> None: + dialect_name = op.get_context().dialect.name + if dialect_name == "postgresql": + op.execute( + f"""UPDATE {table_name} + SET config_snapshot = (config_snapshot::jsonb - 'files')::text + WHERE config_snapshot::jsonb ? 'files'""" + ) + return + if dialect_name == "mysql": + op.execute( + f"""UPDATE {table_name} + SET config_snapshot = JSON_REMOVE(config_snapshot, '$.files') + WHERE JSON_CONTAINS_PATH(config_snapshot, 'one', '$.files')""" + ) + return + if dialect_name == "sqlite": + op.execute( + f"""UPDATE {table_name} + SET config_snapshot = json_remove(config_snapshot, '$.files') + WHERE json_type(config_snapshot, '$.files') IS NOT NULL""" + ) + return + raise RuntimeError(f"unsupported database dialect: {dialect_name}") + + +def upgrade() -> None: + # The Agent Drive removal migration skipped its Python row rewrite when + # migrations were emitted as offline SQL. Run the cleanup again as native + # SQL so every deployment removes the retired AgentSoulConfig.files field. + _remove_legacy_files("agent_config_snapshots") + _remove_legacy_files("agent_config_drafts") + + +def downgrade() -> None: + # The retired files catalog cannot be reconstructed after Agent Drive data + # has been removed. + pass diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index b7643e786b7..07b93004fa8 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -27,7 +27,12 @@ Get account avatar url | ---- | ----------- | ------ | | 200 | Success | **application/json**: [AvatarUrlResponse](#avatarurlresponse)
| -### [POST] /account/avatar +### ~~[POST] /account/avatar~~ + +***DEPRECATED*** + +Deprecated. Use PATCH /account/profile instead. + #### Request Body | Required | Schema | @@ -187,7 +192,12 @@ Get account avatar url | ---- | ----------- | ------ | | 200 | Success | **application/json**: [AccountIntegrateListResponse](#accountintegratelistresponse)
| -### [POST] /account/interface-language +### ~~[POST] /account/interface-language~~ + +***DEPRECATED*** + +Deprecated. Use PATCH /account/profile instead. + #### Request Body | Required | Schema | @@ -200,7 +210,12 @@ Get account avatar url | ---- | ----------- | ------ | | 200 | Success | **application/json**: [AccountResponse](#accountresponse)
| -### [POST] /account/interface-theme +### ~~[POST] /account/interface-theme~~ + +***DEPRECATED*** + +Deprecated. Use PATCH /account/profile instead. + #### Request Body | Required | Schema | @@ -213,7 +228,12 @@ Get account avatar url | ---- | ----------- | ------ | | 200 | Success | **application/json**: [AccountResponse](#accountresponse)
| -### [POST] /account/name +### ~~[POST] /account/name~~ + +***DEPRECATED*** + +Deprecated. Use PATCH /account/profile instead. + #### Request Body | Required | Schema | @@ -246,7 +266,25 @@ Get account avatar url | ---- | ----------- | ------ | | 200 | Success | **application/json**: [AccountResponse](#accountresponse)
| -### [POST] /account/timezone +### [PATCH] /account/profile +#### Request Body + +| Required | Schema | +| -------- | ------ | +| Yes | **application/json**: [AccountProfilePatchPayload](#accountprofilepatchpayload)
| + +#### Responses + +| Code | Description | Schema | +| ---- | ----------- | ------ | +| 200 | Success | **application/json**: [AccountResponse](#accountresponse)
| + +### ~~[POST] /account/timezone~~ + +***DEPRECATED*** + +Deprecated. Use PATCH /account/profile instead. + #### Request Body | Required | Schema | @@ -6526,7 +6564,8 @@ Check if dataset is in use | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [RecommendedAppDetailNullableResponse](#recommendedappdetailnullableresponse)
| +| 200 | Success | **application/json**: [RecommendedAppDetailResponse](#recommendedappdetailresponse)
| +| 404 | Recommended app not found | | ### [GET] /features **Get feature configuration for current tenant** @@ -11070,6 +11109,24 @@ Reset a draft workflow variable to its default value (snippet scope) | 200 | Workflow published successfully | **application/json**: [WorkflowPublishResponse](#workflowpublishresponse)
| | 400 | No draft workflow found | | +### [DELETE] /snippets/{snippet_id}/workflows/{workflow_id} +**Delete a published snippet workflow version** + +#### Parameters + +| Name | Located in | Description | Required | Schema | +| ---- | ---------- | ----------- | -------- | ------ | +| snippet_id | path | Snippet ID | Yes | string (uuid) | +| workflow_id | path | Workflow ID | Yes | string | + +#### Responses + +| Code | Description | +| ---- | ----------- | +| 204 | Workflow deleted successfully | +| 400 | Workflow is in use | +| 404 | Workflow not found | + ### [PATCH] /snippets/{snippet_id}/workflows/{workflow_id} **Update a published snippet workflow version's display metadata** @@ -11185,7 +11242,7 @@ Remove one or more tag bindings from a target. | Name | Located in | Description | Required | Schema | | ---- | ---------- | ----------- | -------- | ------ | | keyword | query | Search keyword | No | string | -| type | query | Tag type filter | No | string,
**Available values:** "", "app", "knowledge", "snippet" | +| type | query | Tag type filter | Yes | string,
**Available values:** "app", "knowledge", "snippet" | #### Responses @@ -11850,6 +11907,7 @@ Export snippet configuration as DSL | ---- | ---------- | ----------- | -------- | ------ | | snippet_id | path | Snippet ID to export | Yes | string (uuid) | | include_secret | query | Whether to include secret variables | No | string,
**Default:** false | +| workflow_id | query | Specific published workflow version to export | No | string | #### Responses @@ -14749,6 +14807,16 @@ Model class for AI model. | password | string | | No | | repeat_new_password | string | | Yes | +#### AccountProfilePatchPayload + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| avatar | string | | No | +| interface_language | string | | No | +| interface_theme | string,
**Available values:** "dark", "light" | *Enum:* `"dark"`, `"light"` | No | +| name | string | | No | +| timezone | string | | No | + #### AccountResponse | Name | Type | Description | Required | @@ -20130,11 +20198,9 @@ How Dify forwards the end-user's identity to an MCP server. #### IncludeSecretQuery -Query parameter for including secret variables in export. - | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| include_secret | string,
**Default:** false | Whether to include secret variables | No | +| include_secret | string,
**Default:** false | Whether to include secret values in the exported DSL | No | #### IndexingEstimate @@ -24903,12 +24969,6 @@ Whitelist scopes accepted by RBAC app and dataset access config APIs. | permission_keys | [ string ] | | No | | updated_at | integer | | Yes | -#### RecommendedAppDetailNullableResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| RecommendedAppDetailNullableResponse | [RecommendedAppDetailResponse](#recommendedappdetailresponse) | | | - #### RecommendedAppDetailResponse | Name | Type | Description | Required | @@ -25558,6 +25618,15 @@ Payload for syncing snippet draft workflow. | hash | string | | No | | input_fields | [ object ] | | No | +#### SnippetExportQuery + +Query parameters for exporting a snippet workflow as DSL. + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| include_secret | string,
**Default:** false | Whether to include secret variables | No | +| workflow_id | string | Specific published workflow version to export | No | + #### SnippetImportPayload Payload for importing snippet from DSL. @@ -25992,7 +26061,7 @@ Non-sensitive bootstrap snapshot exposed before Console or Web authentication. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | keyword | string | Search keyword | No | -| type | string,
**Available values:** "", "app", "knowledge", "snippet" | Tag type filter
*Enum:* `""`, `"app"`, `"knowledge"`, `"snippet"` | No | +| type | string,
**Available values:** "app", "knowledge", "snippet" | Tag type filter
*Enum:* `"app"`, `"knowledge"`, `"snippet"` | Yes | #### TagListResponse diff --git a/api/openapi/markdown/web-openapi.md b/api/openapi/markdown/web-openapi.md index c217552c3df..f080fc833f4 100644 --- a/api/openapi/markdown/web-openapi.md +++ b/api/openapi/markdown/web-openapi.md @@ -847,6 +847,7 @@ Check if user has permission to access a web application. | 400 | Bad Request | | | 401 | Unauthorized | | | 500 | Internal Server Error | | +| 503 | Web App Access Service Unavailable | | ### [POST] /workflows/run **Run workflow** diff --git a/api/pyproject.toml b/api/pyproject.toml index fc84c35fa68..84f26369752 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -6,7 +6,7 @@ requires-python = "~=3.12.0" dependencies = [ # Legacy: mature and widely deployed "bleach>=6.4.0,<7.0.0", - "boto3>=1.43.56,<2.0.0", + "boto3>=1.43.71,<2.0.0", "celery>=5.6.3,<6.0.0", "croniter>=6.2.2,<7.0.0", "dify-agent", @@ -31,7 +31,7 @@ dependencies = [ "flask-migrate>=4.1.0,<5.0.0", "flask-orjson>=2.0.0,<3.0.0", "flask-restx>=1.3.2,<2.0.0", - "google-cloud-aiplatform>=1.160.0,<2.0.0", + "google-cloud-aiplatform>=1.164.0,<2.0.0", "httpx[socks]==0.28.1", "opentelemetry-distro==0.65b0", "opentelemetry-instrumentation-celery==0.65b0", @@ -207,7 +207,7 @@ storage = [ "bce-python-sdk==0.9.76", "cos-python-sdk-v5>=1.9.44,<2.0.0", "esdk-obs-python>=3.26.6,<4.0.0", - "google-cloud-storage>=3.13.0,<4.0.0", + "google-cloud-storage>=3.13.1,<4.0.0", "opendal==0.46.0", "oss2>=2.19.1,<3.0.0", "supabase>=2.31.0,<3.0.0", diff --git a/api/repositories/account_repository.py b/api/repositories/account_repository.py new file mode 100644 index 00000000000..638f47abf44 --- /dev/null +++ b/api/repositories/account_repository.py @@ -0,0 +1,59 @@ +"""SQLAlchemy implementation of the account persistence port.""" + +from typing import override + +from sqlalchemy.orm import Session, sessionmaker + +from models.account import Account +from services.account_ports import AccountRepository +from services.entities.account_entities import AccountProfileChanges, AccountSnapshot + + +class SQLAlchemyAccountRepository(AccountRepository): + def __init__(self, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def get(self, account_id: str) -> AccountSnapshot | None: + with self._session_factory() as session: + account = session.get(Account, account_id) + return self._to_snapshot(account) if account is not None else None + + @override + def update_profile(self, account_id: str, changes: AccountProfileChanges) -> AccountSnapshot | None: + with self._session_factory.begin() as session: + account = session.get(Account, account_id) + if account is None: + return None + + if changes.name is not None: + account.name = changes.name + if changes.avatar is not None: + account.avatar = changes.avatar + if changes.interface_language is not None: + account.interface_language = changes.interface_language + if changes.interface_theme is not None: + account.interface_theme = changes.interface_theme + if changes.timezone is not None: + account.timezone = changes.timezone + + session.flush() + return self._to_snapshot(account) + + @staticmethod + def _to_snapshot(account: Account) -> AccountSnapshot: + return AccountSnapshot( + id=account.id, + name=account.name, + email=account.email, + avatar=account.avatar, + is_password_set=account.is_password_set, + interface_language=account.interface_language, + interface_theme=account.interface_theme, + timezone=account.timezone, + last_login_at=account.last_login_at, + last_login_ip=account.last_login_ip, + status=account.status.value, + initialized_at=account.initialized_at, + created_at=account.created_at, + ) diff --git a/api/repositories/app_definition_query_repository.py b/api/repositories/app_definition_query_repository.py index 45566550600..0907040e787 100644 --- a/api/repositories/app_definition_query_repository.py +++ b/api/repositories/app_definition_query_repository.py @@ -9,6 +9,7 @@ from core.agent.publish_visibility import agent_has_workflow_callable_active_sna from core.app.apps.agent_app.app_feature_projection import merge_agent_app_features from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError +from models.account import Tenant, TenantStatus from models.agent import AgentConfigSnapshot from models.agent_config_entities import AgentSoulConfig from models.model import App, AppMode, AppModelConfig, Site, load_annotation_reply_config @@ -21,6 +22,27 @@ from services.app_definition_query_service import ( AppSiteConfiguration, AppToolIconSource, ) +from services.web_app_runtime_query_service import WebAppRuntimeRecord + + +def _map_site_configuration(site: Site) -> AppSiteConfiguration: + return AppSiteConfiguration( + title=site.title, + chat_color_theme=site.chat_color_theme, + chat_color_theme_inverted=site.chat_color_theme_inverted, + icon_type=site.icon_type.value if site.icon_type is not None else None, + icon=site.icon, + icon_background=site.icon_background, + description=site.description, + copyright=site.copyright, + privacy_policy=site.privacy_policy, + input_placeholder=site.input_placeholder, + custom_disclaimer=site.custom_disclaimer, + default_language=site.default_language, + prompt_public=site.prompt_public, + show_workflow_steps=site.show_workflow_steps, + use_icon_as_answer_icon=site.use_icon_as_answer_icon, + ) def _get_public_agent_parameter_config(app: App, *, session: Session) -> AppParameterConfig: @@ -153,21 +175,41 @@ class AppDefinitionQueryRepository(AppDefinitionQuery): if site is None: return None - return AppSiteConfiguration( - title=site.title, - chat_color_theme=site.chat_color_theme, - chat_color_theme_inverted=site.chat_color_theme_inverted, - icon_type=site.icon_type.value if site.icon_type is not None else None, - icon=site.icon, - icon_background=site.icon_background, - description=site.description, - copyright=site.copyright, - privacy_policy=site.privacy_policy, - input_placeholder=site.input_placeholder, - custom_disclaimer=site.custom_disclaimer, - default_language=site.default_language, - show_workflow_steps=site.show_workflow_steps, - use_icon_as_answer_icon=site.use_icon_as_answer_icon, + return _map_site_configuration(site) + + def get_runtime_record(self, app_id: str) -> WebAppRuntimeRecord | None: + with self._session_factory() as session: + app = session.get(App, app_id) + if app is None: + return None + + site = session.scalar(select(Site).where(Site.app_id == app_id).limit(1)) + if site is None: + return None + + tenant = session.get(Tenant, app.tenant_id) + if tenant is None: + return None + + app_id = app.id + tenant_id = app.tenant_id + enable_site = app.enable_site + site_configuration = _map_site_configuration(site) + plan = tenant.plan + tenant_status = tenant.status.value + tenant_custom_config_json = tenant.custom_config + mode = AppMode.value_of(app.mode).value + if tenant.status != TenantStatus.ARCHIVE: + mode = AppMode.value_of(app.mode_compatible_with_agent_with_session(session=session)).value + return WebAppRuntimeRecord( + app_id=app_id, + tenant_id=tenant_id, + mode=mode, + enable_site=enable_site, + site=site_configuration, + plan=plan, + tenant_status=tenant_status, + tenant_custom_config_json=tenant_custom_config_json, ) @staticmethod diff --git a/api/repositories/recommended_app_catalog_repository.py b/api/repositories/recommended_app_catalog_repository.py new file mode 100644 index 00000000000..9f1cd0dd408 --- /dev/null +++ b/api/repositories/recommended_app_catalog_repository.py @@ -0,0 +1,191 @@ +"""Database-backed recommended app catalog adapter.""" + +import json +import logging +from collections.abc import Sequence +from typing import cast, override + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from constants.languages import languages +from extensions.ext_redis import RedisClientWrapper +from models.model import App, RecommendedApp +from services.app_dsl_service import AppDslService +from services.recommended_app_query_service import ( + RecommendedAppCatalogPage, + RecommendedAppCatalogQuery, + RecommendedAppDetailRecord, + RecommendedAppInfoRecord, + RecommendedAppRecord, +) + +logger = logging.getLogger(__name__) + +# Keep the legacy "explore" Redis key: Explore was the former UI name for this recommended-app surface. +_CATEGORY_ORDER_KEY_PREFIX = "explore:apps:category_order" + + +class DatabaseRecommendedAppCatalogRepository(RecommendedAppCatalogQuery): + def __init__(self, session_factory: sessionmaker[Session], *, redis: RedisClientWrapper) -> None: + self._session_factory = session_factory + self._redis = redis + + @override + def list_recommended(self, language: str) -> RecommendedAppCatalogPage: + with self._session_factory() as session: + recommended_apps = self._list_rows(language, session=session) + if not recommended_apps: + recommended_apps = self._list_rows(languages[0], session=session) + records, categories = self._map_rows(recommended_apps, session=session) + return RecommendedAppCatalogPage( + recommended_apps=records, + categories=tuple(self._order_categories(categories, language)), + ) + + @override + def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: + with self._session_factory() as session: + recommended_apps = self._list_rows(language, session=session, is_learn_dify=True) + if not recommended_apps and language != languages[0]: + recommended_apps = self._list_rows(languages[0], session=session, is_learn_dify=True) + records, _ = self._map_rows(recommended_apps, session=session) + return RecommendedAppCatalogPage(recommended_apps=records, categories=()) + + @override + def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: + with self._session_factory() as session: + return self._get_detail(app_id, session=session) + + @override + def contains(self, app_id: str) -> bool: + with self._session_factory() as session: + return ( + session.scalar( + select(RecommendedApp.app_id) + .join(App, App.id == RecommendedApp.app_id) + .where( + RecommendedApp.app_id == app_id, + RecommendedApp.is_listed.is_(True), + App.is_public.is_(True), + ) + .limit(1) + ) + is not None + ) + + def _order_categories(self, categories: set[str], language: str) -> list[str]: + try: + raw_categories = self._redis.get(f"{_CATEGORY_ORDER_KEY_PREFIX}:{language}") + except Exception: + logger.exception("Failed to read recommended app category order from Redis.") + return sorted(categories) + + if not raw_categories: + return sorted(categories) + if isinstance(raw_categories, bytes): + raw_categories = raw_categories.decode("utf-8") + + try: + configured_order = json.loads(raw_categories) + except (TypeError, json.JSONDecodeError): + logger.warning("Invalid recommended app category order payload for language %s.", language) + return sorted(categories) + + if not isinstance(configured_order, list): + return sorted(categories) + + string_order = [category for category in configured_order if isinstance(category, str)] + return string_order or sorted(categories) + + @staticmethod + def _list_rows( + language: str, + *, + session: Session, + is_learn_dify: bool | None = None, + ) -> list[RecommendedApp]: + filters = [RecommendedApp.is_listed.is_(True), RecommendedApp.language == language] + if is_learn_dify is not None: + filters.append(RecommendedApp.is_learn_dify.is_(is_learn_dify)) + return list(session.scalars(select(RecommendedApp).where(*filters)).all()) + + @classmethod + def _map_rows( + cls, + recommended_apps: Sequence[RecommendedApp], + *, + session: Session, + ) -> tuple[tuple[RecommendedAppRecord, ...], set[str]]: + categories: set[str] = set() + records: list[RecommendedAppRecord] = [] + for recommended_app in recommended_apps: + app = session.get(App, recommended_app.app_id) + if app is None or not app.is_public: + continue + + site = app.site_with_session(session=session) + if site is None: + continue + + app_categories = cls._as_string_tuple(recommended_app.categories or (), field="categories") + records.append( + RecommendedAppRecord( + app=RecommendedAppInfoRecord( + id=app.id, + name=app.name, + mode=app.mode.value, + icon=cast(str | None, app.icon), + icon_type=app.icon_type.value if app.icon_type is not None else None, + icon_background=app.icon_background, + ), + app_id=recommended_app.app_id, + description=cast(str | None, site.description), + copyright=cast(str | None, site.copyright), + privacy_policy=cast(str | None, site.privacy_policy), + custom_disclaimer=cast(str | None, site.custom_disclaimer), + categories=app_categories, + position=recommended_app.position, + is_listed=recommended_app.is_listed, + ) + ) + categories.update(app_categories) + + return tuple(records), categories + + @staticmethod + def _get_detail(app_id: str, *, session: Session) -> RecommendedAppDetailRecord | None: + recommended_app = session.scalar( + select(RecommendedApp) + .where( + RecommendedApp.is_listed.is_(True), + RecommendedApp.app_id == app_id, + ) + .limit(1) + ) + if recommended_app is None: + return None + + app = session.get(App, app_id) + if app is None or not app.is_public: + return None + + return RecommendedAppDetailRecord( + id=app.id, + name=app.name, + icon=cast(str | None, app.icon), + icon_background=app.icon_background, + mode=app.mode.value, + export_data=AppDslService.export_dsl(app_model=app, session=session), + ) + + @staticmethod + def _as_string_tuple(value: object, *, field: str) -> tuple[str, ...]: + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError(f"{field} must be a sequence of strings") + items: list[str] = [] + for item in value: + if not isinstance(item, str): + raise TypeError(f"{field} must contain only strings") + items.append(item) + return tuple(items) diff --git a/api/repositories/tag_repository.py b/api/repositories/tag_repository.py new file mode 100644 index 00000000000..9b949ebc6f0 --- /dev/null +++ b/api/repositories/tag_repository.py @@ -0,0 +1,214 @@ +"""SQLAlchemy persistence adapter for Console tag management.""" + +import uuid +from typing import override + +import sqlalchemy as sa +from sqlalchemy import delete, func, select +from sqlalchemy.orm import Session, sessionmaker + +from libs.helper import escape_like_pattern +from models.dataset import Dataset +from models.enums import TagType +from models.model import App, Tag, TagBinding +from models.snippet import CustomizedSnippet +from services.tag_application_service import ( + CreateTagInput, + InvalidTagBindingTypeError, + TagBindingInput, + TagBindingTargetNotFoundError, + TagNameConflictError, + TagNotFoundError, + TagStore, + TagSummary, + UpdateTagInput, +) + + +class TagRepository(TagStore): + def __init__(self, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def list_tags(self, workspace_id: str, tag_type: str, keyword: str | None) -> tuple[TagSummary, ...]: + stmt = ( + select(Tag.id, Tag.name, Tag.type, func.count(TagBinding.id)) + .outerjoin( + TagBinding, + sa.and_(TagBinding.tag_id == Tag.id, TagBinding.tenant_id == workspace_id), + ) + .where(Tag.type == tag_type, Tag.tenant_id == workspace_id) + ) + if keyword: + escaped_keyword = escape_like_pattern(keyword) + stmt = stmt.where(Tag.name.ilike(f"%{escaped_keyword}%", escape="\\")) + stmt = stmt.group_by(Tag.id, Tag.name, Tag.type, Tag.created_at).order_by(Tag.created_at.desc()) + + with self._session_factory() as session: + return tuple( + TagSummary( + id=tag_id, + name=name, + type=tag_kind.value, + binding_count=binding_count, + ) + for tag_id, name, tag_kind, binding_count in session.execute(stmt).all() + ) + + @override + def get_tag_type(self, workspace_id: str, tag_id: str) -> str | None: + with self._session_factory() as session: + tag_type = session.scalar(select(Tag.type).where(Tag.id == tag_id, Tag.tenant_id == workspace_id).limit(1)) + return tag_type.value if tag_type is not None else None + + @override + def create_tag(self, workspace_id: str, actor_id: str, tag: CreateTagInput) -> TagSummary: + with self._session_factory.begin() as session: + existing = session.scalar( + select(Tag.id).where(Tag.name == tag.name, Tag.tenant_id == workspace_id, Tag.type == tag.type).limit(1) + ) + if existing is not None: + raise TagNameConflictError + + model = Tag( + name=tag.name, + type=TagType(tag.type), + created_by=actor_id, + tenant_id=workspace_id, + ) + model.id = str(uuid.uuid4()) + session.add(model) + session.flush() + return self._summary(model, binding_count=0) + + @override + def update_tag(self, workspace_id: str, tag_id: str, tag: UpdateTagInput) -> TagSummary: + with self._session_factory.begin() as session: + model = session.scalar(select(Tag).where(Tag.id == tag_id, Tag.tenant_id == workspace_id).limit(1)) + if model is None: + raise TagNotFoundError + + if tag.name != model.name: + existing = session.scalar( + select(Tag.id) + .where( + Tag.name == tag.name, + Tag.tenant_id == workspace_id, + Tag.type == model.type, + Tag.id != tag_id, + ) + .limit(1) + ) + if existing is not None: + raise TagNameConflictError + model.name = tag.name + + binding_count = ( + session.scalar( + select(func.count(TagBinding.id)).where( + TagBinding.tag_id == tag_id, + TagBinding.tenant_id == workspace_id, + ) + ) + or 0 + ) + return self._summary(model, binding_count=binding_count) + + @override + def delete_tag(self, workspace_id: str, tag_id: str) -> None: + with self._session_factory.begin() as session: + model = session.scalar(select(Tag).where(Tag.id == tag_id, Tag.tenant_id == workspace_id).limit(1)) + if model is None: + raise TagNotFoundError + + session.execute( + delete(TagBinding).where( + TagBinding.tag_id == tag_id, + TagBinding.tenant_id == workspace_id, + ) + ) + session.delete(model) + + @override + def create_bindings(self, workspace_id: str, actor_id: str, binding: TagBindingInput) -> None: + with self._session_factory.begin() as session: + self._ensure_target_exists(session, workspace_id, binding) + requested_tag_ids = tuple(dict.fromkeys(binding.tag_ids)) + if not requested_tag_ids: + return + + valid_tag_ids = tuple( + session.scalars( + select(Tag.id).where( + Tag.id.in_(requested_tag_ids), + Tag.tenant_id == workspace_id, + Tag.type == binding.type, + ) + ).all() + ) + if not valid_tag_ids: + return + + existing_tag_ids = set( + session.scalars( + select(TagBinding.tag_id).where( + TagBinding.tag_id.in_(valid_tag_ids), + TagBinding.target_id == binding.target_id, + TagBinding.tenant_id == workspace_id, + ) + ).all() + ) + session.add_all( + TagBinding( + tag_id=tag_id, + target_id=binding.target_id, + tenant_id=workspace_id, + created_by=actor_id, + ) + for tag_id in valid_tag_ids + if tag_id not in existing_tag_ids + ) + + @override + def delete_bindings(self, workspace_id: str, binding: TagBindingInput) -> None: + with self._session_factory.begin() as session: + self._ensure_target_exists(session, workspace_id, binding) + session.execute( + delete(TagBinding).where( + TagBinding.target_id == binding.target_id, + TagBinding.tag_id.in_(binding.tag_ids), + TagBinding.tenant_id == workspace_id, + TagBinding.tag_id.in_( + select(Tag.id).where( + Tag.tenant_id == workspace_id, + Tag.type == binding.type, + ) + ), + ) + ) + + @staticmethod + def _summary(tag: Tag, *, binding_count: int) -> TagSummary: + return TagSummary( + id=tag.id, + name=tag.name, + type=tag.type.value, + binding_count=binding_count, + ) + + @staticmethod + def _ensure_target_exists(session: Session, workspace_id: str, binding: TagBindingInput) -> None: + if binding.type == "knowledge": + stmt = select(Dataset.id).where(Dataset.tenant_id == workspace_id, Dataset.id == binding.target_id) + elif binding.type == "app": + stmt = select(App.id).where(App.tenant_id == workspace_id, App.id == binding.target_id) + elif binding.type == "snippet": + stmt = select(CustomizedSnippet.id).where( + CustomizedSnippet.tenant_id == workspace_id, + CustomizedSnippet.id == binding.target_id, + ) + else: + raise InvalidTagBindingTypeError + + if session.scalar(stmt.limit(1)) is None: + raise TagBindingTargetNotFoundError(binding.type) diff --git a/api/repositories/trial_app_query_repository.py b/api/repositories/trial_app_query_repository.py new file mode 100644 index 00000000000..d8890724466 --- /dev/null +++ b/api/repositories/trial_app_query_repository.py @@ -0,0 +1,23 @@ +"""Database repository for recommended app trial eligibility.""" + +from collections.abc import Sequence, Set +from typing import override + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from models.model import TrialApp +from services.recommended_app_query_service import TrialAppQuery + + +class TrialAppQueryRepository(TrialAppQuery): + def __init__(self, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def existing_ids(self, app_ids: Sequence[str]) -> Set[str]: + if not app_ids: + return frozenset() + + with self._session_factory() as session: + return frozenset(session.scalars(select(TrialApp.app_id).where(TrialApp.app_id.in_(app_ids))).all()) diff --git a/api/repositories/trial_app_usage_repository.py b/api/repositories/trial_app_usage_repository.py new file mode 100644 index 00000000000..a9c60738445 --- /dev/null +++ b/api/repositories/trial_app_usage_repository.py @@ -0,0 +1,28 @@ +"""Database repository for recommended trial app usage.""" + +from typing import override + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from models.model import AccountTrialAppRecord +from services.trial_app_usage import TrialAppUsageRecorder + + +class TrialAppUsageRepository(TrialAppUsageRecorder): + def __init__(self, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def record(self, *, app_id: str, account_id: str) -> None: + """Increment usage without committing the caller's request transaction.""" + with self._session_factory() as session, session.begin(): + record = session.scalar( + select(AccountTrialAppRecord) + .where(AccountTrialAppRecord.app_id == app_id, AccountTrialAppRecord.account_id == account_id) + .limit(1) + ) + if record is None: + session.add(AccountTrialAppRecord(app_id=app_id, account_id=account_id, count=1)) + else: + record.count += 1 diff --git a/api/services/account_activation_adapters.py b/api/services/account_activation_adapters.py index ca22e7983b8..62b58c1336c 100644 --- a/api/services/account_activation_adapters.py +++ b/api/services/account_activation_adapters.py @@ -54,8 +54,10 @@ class BillingAccountActivationEligibility(AccountActivationEligibility): self._enabled = enabled @override - def is_frozen(self, email: str) -> bool: - return self._enabled and BillingService.is_email_in_freeze(email) + def get_freeze_type(self, email: str) -> str | None: + if not self._enabled: + return None + return BillingService.get_email_freeze_type(email) class BillingWorkspaceMembershipCache(WorkspaceMembershipCache): diff --git a/api/services/account_activation_service.py b/api/services/account_activation_service.py index 1aaafab538d..d635cf24f3d 100644 --- a/api/services/account_activation_service.py +++ b/api/services/account_activation_service.py @@ -41,7 +41,7 @@ class WorkspaceInvitePolicy(Protocol): class AccountActivationEligibility(Protocol): - def is_frozen(self, email: str) -> bool: ... + def get_freeze_type(self, email: str) -> str | None: ... class WorkspaceMembershipCache(Protocol): @@ -60,6 +60,10 @@ class FrozenAccountError(Exception): """The invited account is temporarily ineligible for activation.""" +class EmailDomainSuspendedError(Exception): + """The invited account uses a suspended email domain.""" + + class AccountActivationService: def __init__( self, @@ -101,7 +105,10 @@ class AccountActivationService: if authenticated_account_id is not None and authenticated_account_id != invitation.account_id: raise InvitationAccountMismatchError - if self._eligibility.is_frozen(invitation.account_email): + freeze_type = self._eligibility.get_freeze_type(invitation.account_email) + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedError + if freeze_type: raise FrozenAccountError setup = self._resolve_setup(invitation, command) diff --git a/api/services/account_errors.py b/api/services/account_errors.py new file mode 100644 index 00000000000..2aba4cd4be7 --- /dev/null +++ b/api/services/account_errors.py @@ -0,0 +1,9 @@ +"""Framework-neutral errors shared by account application services.""" + + +class AccountApplicationError(Exception): + """Base class for failures owned by account application services.""" + + +class AccountNotFoundError(AccountApplicationError): + """The admitted account no longer exists.""" diff --git a/api/services/account_ports.py b/api/services/account_ports.py new file mode 100644 index 00000000000..6de05706119 --- /dev/null +++ b/api/services/account_ports.py @@ -0,0 +1,11 @@ +"""Persistence ports used by account application services.""" + +from typing import Protocol + +from services.entities.account_entities import AccountProfileChanges, AccountSnapshot + + +class AccountRepository(Protocol): + def get(self, account_id: str) -> AccountSnapshot | None: ... + + def update_profile(self, account_id: str, changes: AccountProfileChanges) -> AccountSnapshot | None: ... diff --git a/api/services/account_profile_service.py b/api/services/account_profile_service.py new file mode 100644 index 00000000000..7c492567b8e --- /dev/null +++ b/api/services/account_profile_service.py @@ -0,0 +1,26 @@ +"""Application service for reading and updating the current account profile.""" + +from machinery.context import RequestContext +from services.account_errors import AccountNotFoundError +from services.account_ports import AccountRepository +from services.entities.account_entities import AccountProfileChanges, AccountSnapshot + + +class AccountProfileService: + def __init__(self, *, accounts: AccountRepository) -> None: + self._accounts = accounts + + def get(self, context: RequestContext) -> AccountSnapshot: + account = self._accounts.get(context.account_id) + if account is None: + raise AccountNotFoundError + return account + + def update(self, context: RequestContext, changes: AccountProfileChanges) -> AccountSnapshot: + if changes.has_changes(): + account = self._accounts.update_profile(context.account_id, changes) + else: + account = self._accounts.get(context.account_id) + if account is None: + raise AccountNotFoundError + return account diff --git a/api/services/account_service.py b/api/services/account_service.py index 367493cd5e2..fcc9a3e90b9 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -67,6 +67,7 @@ from services.errors.account import ( AccountRegisterError, CannotOperateSelfError, CurrentPasswordIncorrectError, + EmailDomainSuspendedError, InvalidActionError, LinkAccountIntegrateError, MemberNotInTenantError, @@ -470,6 +471,9 @@ class AccountService: raise SeatsLimitExceededError("licensed seats limit exceeded") if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and BillingService.is_email_in_freeze(email): + freeze_type = BillingService.get_email_freeze_type(email) or "freeze" + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedError() raise AccountRegisterError( description=( "This email account has been deleted within the past " @@ -1070,6 +1074,9 @@ class AccountService: @classmethod def get_user_through_email(cls, email: str, *, session: Session): if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and BillingService.is_email_in_freeze(email): + freeze_type = BillingService.get_email_freeze_type(email) or "freeze" + if freeze_type == "email_domain_suspended": + raise EmailDomainSuspendedError() raise AccountRegisterError( description=( "This email account has been deleted within the past " @@ -1088,9 +1095,13 @@ class AccountService: @classmethod def is_account_in_freeze(cls, email: str) -> bool: - if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and BillingService.is_email_in_freeze(email): - return True - return False + return cls.get_account_freeze_type(email) is not None + + @classmethod + def get_account_freeze_type(cls, email: str): + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: + return None + return BillingService.get_email_freeze_type(email) @staticmethod @redis_fallback(default_return=None) diff --git a/api/services/agent/composer_service.py b/api/services/agent/composer_service.py index 064a6dc8e56..5ca7d670ad6 100644 --- a/api/services/agent/composer_service.py +++ b/api/services/agent/composer_service.py @@ -279,16 +279,11 @@ class AgentComposerService: ) state["validation"] = cls.collect_validation_findings(payload=payload) session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id=tenant_id, agent_ids=retirement_candidates, account_id=account_id, ) - enqueue_agent_resource_collection( - tenant_id=tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) return state @classmethod diff --git a/api/services/agent/deletion_service.py b/api/services/agent/deletion_service.py new file mode 100644 index 00000000000..0801b8bb0e7 --- /dev/null +++ b/api/services/agent/deletion_service.py @@ -0,0 +1,103 @@ +"""Hard-delete archived Agent aggregates after external resources are collected.""" + +from __future__ import annotations + +from collections.abc import Iterable + +from sqlalchemy import delete, select + +from core.db.session_factory import session_factory +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigRevision, + AgentConfigSnapshot, + AgentDebugConversation, + AgentHomeSnapshot, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, +) + + +class AgentDeletionInvariantError(RuntimeError): + """An archived Agent no longer satisfies the hard-deletion contract.""" + + +class AgentDeletionService: + """Delete archived Agent aggregates after their external resources are gone. + + The aggregate includes Agent-owned configuration, debug, Home, and Workspace + Binding rows. Workflow-owned binding soft references are outside the + aggregate and may remain dangling after deletion. + """ + + @classmethod + def purge_archived_agents(cls, *, tenant_id: str, agent_ids: Iterable[str]) -> None: + """Idempotently hard-delete eligible archived Agent aggregates. + + Missing targets are a no-op. Every stored target must be ``ARCHIVED``, + have no ACTIVE Workspace Binding or Home Snapshot, and all dependent rows + and Agents are deleted and committed in one transaction; an exception + before commit leaves the transaction to roll back without a partial + aggregate deletion. + """ + candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) + if not candidates: + return + + with session_factory.create_session() as session: + agents = session.scalars(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id.in_(candidates))).all() + if not agents: + return + + stored_ids = [agent.id for agent in agents] + non_archived_ids = [agent.id for agent in agents if agent.status != AgentStatus.ARCHIVED] + if non_archived_ids: + raise AgentDeletionInvariantError( + f"Agents must be ARCHIVED before deletion: {', '.join(non_archived_ids)}" + ) + + active_binding_id = session.scalar( + select(AgentWorkspaceBinding.id) + .where( + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.agent_id.in_(stored_ids), + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + .limit(1) + ) + if active_binding_id is not None: + raise AgentDeletionInvariantError(f"Agent aggregate still has ACTIVE Binding {active_binding_id}") + + active_home_id = session.scalar( + select(AgentHomeSnapshot.id) + .where( + AgentHomeSnapshot.tenant_id == tenant_id, + AgentHomeSnapshot.agent_id.in_(stored_ids), + AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE, + ) + .limit(1) + ) + if active_home_id is not None: + raise AgentDeletionInvariantError(f"Agent aggregate still has ACTIVE Home Snapshot {active_home_id}") + + for model in ( + AgentDebugConversation, + AgentConfigRevision, + AgentConfigDraft, + AgentConfigSnapshot, + AgentHomeSnapshot, + AgentWorkspaceBinding, + ): + session.execute( + delete(model).where( + model.tenant_id == tenant_id, + model.agent_id.in_(stored_ids), + ) + ) + session.execute(delete(Agent).where(Agent.tenant_id == tenant_id, Agent.id.in_(stored_ids))) + session.commit() + + +__all__ = ["AgentDeletionInvariantError", "AgentDeletionService"] diff --git a/api/services/agent/home_snapshot_service.py b/api/services/agent/home_snapshot_service.py index 54b0eb2591e..9e2113c56a6 100644 --- a/api/services/agent/home_snapshot_service.py +++ b/api/services/agent/home_snapshot_service.py @@ -15,7 +15,6 @@ from libs.uuid_utils import uuidv7 from models.agent import ( Agent, AgentConfigDraft, - AgentConfigSnapshot, AgentConfigVersionKind, AgentHomeSnapshot, AgentStatus, @@ -108,13 +107,13 @@ class AgentHomeSnapshotService: select(AgentHomeSnapshot).where( AgentHomeSnapshot.tenant_id == tenant_id, AgentHomeSnapshot.agent_id == agent_id, - AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE, ) ).all() now = naive_utc_now() for row in rows: - row.status = AgentWorkingResourceStatus.RETIRED - row.retired_at = now + if row.status == AgentWorkingResourceStatus.ACTIVE: + row.status = AgentWorkingResourceStatus.RETIRED + row.retired_at = now return [row.id for row in rows] @classmethod @@ -129,13 +128,6 @@ class AgentHomeSnapshotService: ) if snapshot is None: return - referenced = session.scalar( - select(AgentConfigDraft.id).where(AgentConfigDraft.home_snapshot_id == home_snapshot_id).limit(1) - ) or session.scalar( - select(AgentConfigSnapshot.id).where(AgentConfigSnapshot.home_snapshot_id == home_snapshot_id).limit(1) - ) - if referenced is not None: - return snapshot_ref = snapshot.snapshot_ref cls.delete(snapshot_ref=snapshot_ref) with session_factory.create_session() as session: diff --git a/api/services/agent/retirement_service.py b/api/services/agent/retirement_service.py index 25b133e14f9..c800a8ea51b 100644 --- a/api/services/agent/retirement_service.py +++ b/api/services/agent/retirement_service.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging from collections.abc import Iterable -from sqlalchemy import or_, select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from core.db.session_factory import session_factory @@ -15,20 +15,22 @@ from models.agent import ( AgentScope, AgentStatus, AgentWorkingResourceStatus, + AgentWorkspace, AgentWorkspaceBinding, WorkflowAgentNodeBinding, ) -from models.enums import AppStatus -from models.model import App +from models.model import App, AppMode from models.workflow import Workflow from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.workspace_service import AgentWorkspaceService +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection +from tasks.remove_app_and_related_data_task import remove_app_and_related_data_task logger = logging.getLogger(__name__) class WorkflowAgentRetirementService: - """Archive workflow-only Agents once no effective binding owns them.""" + """Delete workflow-only Agent aggregates after their last Workflow owner is gone.""" @classmethod def retire_unowned( @@ -37,13 +39,25 @@ class WorkflowAgentRetirementService: tenant_id: str, agent_ids: Iterable[str], account_id: str | None, - ) -> tuple[list[str], list[str]]: - """Re-check ownership, archive orphans, and commit their resource retirement.""" + ) -> None: + """Retire unowned workflow-only Agents in an independent transaction. + + This method returns ``None``. It archives orphan Agents, retires their + working resources, and deletes their hidden Apps before committing. It + then publishes every hidden-App cleanup before publishing the Agent + resource collector; database and task-publication errors propagate. + + Archived Agents, missing hidden App rows, and already-retired resources + remain cleanup candidates. A retry can therefore publish duplicate + cleanup tasks, which are expected to be idempotent. + """ candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) if not candidates: - return [], [] + return + backing_app_ids: list[str] = [] retired_bindings: list[str] = [] + retired_workspaces: list[str] = [] retired_snapshots: list[str] = [] try: with session_factory.create_session() as session: @@ -53,22 +67,43 @@ class WorkflowAgentRetirementService: agent_ids=candidates, account_id=account_id, ) + retired_agents = session.scalars( + select(Agent).where( + Agent.tenant_id == tenant_id, + Agent.id.in_(retired_agent_ids), + ) + ).all() + backing_app_ids = sorted({agent.backing_app_id for agent in retired_agents if agent.backing_app_id}) + for app_id in backing_app_ids: + AgentWorkspaceService.retire_all_for_app( + session=session, + tenant_id=tenant_id, + app_id=app_id, + ) + retired_workspaces.extend( + session.scalars( + select(AgentWorkspace.id).where( + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.app_id == app_id, + AgentWorkspace.status == AgentWorkingResourceStatus.RETIRED, + ) + ).all() + ) for agent_id in retired_agent_ids: bindings = session.scalars( select(AgentWorkspaceBinding).where( AgentWorkspaceBinding.tenant_id == tenant_id, AgentWorkspaceBinding.agent_id == agent_id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, ) ).all() for binding in bindings: - binding_id = AgentWorkspaceService.retire_binding( - session=session, - tenant_id=tenant_id, - binding_id=binding.id, - ) - if binding_id is not None: - retired_bindings.append(binding_id) + if binding.status == AgentWorkingResourceStatus.ACTIVE: + AgentWorkspaceService.retire_binding( + session=session, + tenant_id=tenant_id, + binding_id=binding.id, + ) + retired_bindings.append(binding.id) retired_snapshots.extend( AgentHomeSnapshotService.retire_all_for_agent( session=session, @@ -76,6 +111,14 @@ class WorkflowAgentRetirementService: agent_id=agent_id, ) ) + if backing_app_ids: + session.execute( + delete(App).where( + App.tenant_id == tenant_id, + App.id.in_(backing_app_ids), + App.mode == AppMode.AGENT, + ) + ) session.commit() except Exception: logger.exception( @@ -85,8 +128,24 @@ class WorkflowAgentRetirementService: "agent_ids": candidates, }, ) - return [], [] - return retired_bindings, retired_snapshots + raise + + for app_id in backing_app_ids: + try: + remove_app_and_related_data_task.delay(tenant_id=tenant_id, app_id=app_id) + except Exception: + logger.exception( + "Failed to enqueue hidden Agent App cleanup", + extra={"tenant_id": tenant_id, "app_id": app_id}, + ) + raise + enqueue_agent_resource_collection( + tenant_id=tenant_id, + workspace_ids=retired_workspaces, + binding_ids=retired_bindings, + home_snapshot_ids=retired_snapshots, + purge_agent_ids=retired_agent_ids, + ) @classmethod def archive_unowned( @@ -97,7 +156,7 @@ class WorkflowAgentRetirementService: agent_ids: Iterable[str], account_id: str | None, ) -> list[str]: - """Archive active orphans and return every orphan eligible for Home cleanup.""" + """Archive active orphans and return complete aggregate purge candidates.""" candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) if not candidates: return [] @@ -109,7 +168,7 @@ class WorkflowAgentRetirementService: Agent.status.in_((AgentStatus.ACTIVE, AgentStatus.ARCHIVED)), ) ).all() - effective_agent_ids = cls._effective_agent_ids( + retained_agent_ids = cls.retained_agent_ids( session=session, tenant_id=tenant_id, agent_ids=[agent.id for agent in agents], @@ -117,7 +176,7 @@ class WorkflowAgentRetirementService: now = naive_utc_now() cleanup_candidates: list[str] = [] for agent in agents: - if agent.id in effective_agent_ids: + if agent.id in retained_agent_ids: continue if agent.status == AgentStatus.ACTIVE: agent.status = AgentStatus.ARCHIVED @@ -130,33 +189,32 @@ class WorkflowAgentRetirementService: return cleanup_candidates @staticmethod - def _effective_agent_ids( + def retained_agent_ids( *, session: Session, tenant_id: str, agent_ids: list[str], ) -> set[str]: + """Return Agents that still have an exact persisted Workflow owner. + + The owner key is tenant, App, Workflow, and Workflow version. Draft and + every published version, whether current or historical, count equally; + the App's current-Workflow pointer is not part of ownership. + """ if not agent_ids: return set() values = session.scalars( select(WorkflowAgentNodeBinding.agent_id) .join( Workflow, - Workflow.id == WorkflowAgentNodeBinding.workflow_id, + (Workflow.tenant_id == WorkflowAgentNodeBinding.tenant_id) + & (Workflow.app_id == WorkflowAgentNodeBinding.app_id) + & (Workflow.id == WorkflowAgentNodeBinding.workflow_id) + & (Workflow.version == WorkflowAgentNodeBinding.workflow_version), ) - .join(App, App.id == WorkflowAgentNodeBinding.app_id) .where( WorkflowAgentNodeBinding.tenant_id == tenant_id, WorkflowAgentNodeBinding.agent_id.in_(agent_ids), - Workflow.tenant_id == tenant_id, - Workflow.app_id == WorkflowAgentNodeBinding.app_id, - Workflow.version == WorkflowAgentNodeBinding.workflow_version, - App.tenant_id == tenant_id, - App.status == AppStatus.NORMAL, - or_( - Workflow.version == Workflow.VERSION_DRAFT, - App.workflow_id == Workflow.id, - ), ) .distinct() ).all() diff --git a/api/services/agent/roster_service.py b/api/services/agent/roster_service.py index af7bed03879..d63056209ef 100644 --- a/api/services/agent/roster_service.py +++ b/api/services/agent/roster_service.py @@ -23,8 +23,6 @@ from models.agent import ( AgentScope, AgentSource, AgentStatus, - AgentWorkingResourceStatus, - AgentWorkspaceBinding, AgentWorkspaceOwnerType, WorkflowAgentBindingType, WorkflowAgentNodeBinding, @@ -42,7 +40,6 @@ from services.agent.errors import ( AgentNotFoundError, AgentVersionNotFoundError, ) -from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope from services.app_service import AppService, CreateAppParams from services.enterprise.enterprise_service import EnterpriseService @@ -1261,41 +1258,6 @@ class AgentRosterService: raise AgentNameConflictError() from exc return self.get_roster_agent_detail(tenant_id=tenant_id, agent_id=agent_id) - def archive_roster_agent(self, *, tenant_id: str, agent_id: str, account_id: str) -> None: - agent = self._get_agent(tenant_id=tenant_id, agent_id=agent_id, roster_only=True) - retired_binding_ids: list[str] = [] - if agent.status != AgentStatus.ARCHIVED: - agent.status = AgentStatus.ARCHIVED - agent.archived_by = account_id - agent.archived_at = naive_utc_now() - agent.updated_by = account_id - bindings = self._session.scalars( - select(AgentWorkspaceBinding).where( - AgentWorkspaceBinding.tenant_id == tenant_id, - AgentWorkspaceBinding.agent_id == agent_id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, - ) - ).all() - for binding in bindings: - retired_id = AgentWorkspaceService.retire_binding( - session=self._session, - tenant_id=tenant_id, - binding_id=binding.id, - ) - if retired_id is not None: - retired_binding_ids.append(retired_id) - retired_snapshot_ids = AgentHomeSnapshotService.retire_all_for_agent( - session=self._session, - tenant_id=tenant_id, - agent_id=agent_id, - ) - self._session.commit() - enqueue_agent_resource_collection( - tenant_id=tenant_id, - binding_ids=retired_binding_ids, - home_snapshot_ids=retired_snapshot_ids, - ) - @staticmethod def _visible_version_operations(agent: Agent) -> set[AgentConfigRevisionOperation]: if agent.source == AgentSource.AGENT_APP or ( diff --git a/api/services/agent/workflow_publish_service.py b/api/services/agent/workflow_publish_service.py index 954a3819774..117d8f12ef0 100644 --- a/api/services/agent/workflow_publish_service.py +++ b/api/services/agent/workflow_publish_service.py @@ -24,7 +24,6 @@ from models.agent_config_entities import ( WorkflowNodeJobConfig, WorkflowPreviousNodeOutputRef, ) -from models.model import App from models.workflow import Workflow from services.agent.composer_validator import ComposerConfigValidator from services.agent.prompt_mentions import ( @@ -575,32 +574,18 @@ class WorkflowAgentPublishService: session: Session, draft_workflow: Workflow, published_workflow: Workflow, - ) -> set[str]: - current_workflow_id = session.scalar( - select(App.workflow_id).where( - App.tenant_id == draft_workflow.tenant_id, - App.id == draft_workflow.app_id, - ) - ) - retirement_candidates: set[str] = set() - if current_workflow_id: - retirement_candidates = { - agent_id - for agent_id in session.scalars( - select(WorkflowAgentNodeBinding.agent_id).where( - WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id, - WorkflowAgentNodeBinding.app_id == draft_workflow.app_id, - WorkflowAgentNodeBinding.workflow_id == current_workflow_id, - WorkflowAgentNodeBinding.binding_type == WorkflowAgentBindingType.INLINE_AGENT, - ) - ).all() - if agent_id - } + ) -> None: + """Copy all draft Roster and inline bindings to a published version. + + Only copied inline bindings add owners for workflow-only Agents. + Publishing does not release existing draft or historical inline owners, + produces no retirement candidates, and returns ``None``. + """ node_ids = { node_id for node_id, _node_data in WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict) } if not node_ids: - return retirement_candidates + return bindings = session.scalars( select(WorkflowAgentNodeBinding).where( @@ -612,25 +597,17 @@ class WorkflowAgentPublishService: ) ).all() if not bindings: - return retirement_candidates - - agents_by_id = { - agent.id: agent - for agent in session.scalars( - select(Agent).where( - Agent.tenant_id == draft_workflow.tenant_id, - Agent.id.in_({binding.agent_id for binding in bindings if binding.agent_id}), - ) - ).all() - } + return for binding in bindings: - agent = agents_by_id.get(binding.agent_id) if binding.agent_id else None - current_snapshot_id = ( - agent.active_config_snapshot_id - if agent is not None and binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT - else binding.current_snapshot_id - ) + current_snapshot_id = binding.current_snapshot_id + if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT and binding.agent_id: + _, current_snapshot_id = cls._resolve_roster_agent_graph_binding( + session=session, + draft_workflow=draft_workflow, + node_id=binding.node_id, + agent_id=binding.agent_id, + ) copied = WorkflowAgentNodeBinding( tenant_id=binding.tenant_id, app_id=binding.app_id, @@ -645,7 +622,6 @@ class WorkflowAgentPublishService: updated_by=binding.updated_by, ) session.add(copied) - return retirement_candidates @classmethod def restore_agent_node_bindings_to_draft( @@ -671,9 +647,6 @@ class WorkflowAgentPublishService: for binding in existing if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id } - for binding in existing: - session.delete(binding) - source_bindings = session.scalars( select(WorkflowAgentNodeBinding).where( WorkflowAgentNodeBinding.tenant_id == source_workflow.tenant_id, @@ -682,6 +655,19 @@ class WorkflowAgentPublishService: WorkflowAgentNodeBinding.workflow_version == source_workflow.version, ) ).all() + for source in source_bindings: + if source.binding_type == WorkflowAgentBindingType.ROSTER_AGENT and source.agent_id: + cls._resolve_roster_agent_graph_binding( + session=session, + draft_workflow=draft_workflow, + node_id=source.node_id, + agent_id=source.agent_id, + ) + + for binding in existing: + session.delete(binding) + session.flush() + for source in source_bindings: agent_id = source.agent_id snapshot_id = source.current_snapshot_id diff --git a/api/services/agent/workspace_service.py b/api/services/agent/workspace_service.py index 2cc6c34e881..cdf2736b523 100644 --- a/api/services/agent/workspace_service.py +++ b/api/services/agent/workspace_service.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from dify_agent.client import Client from dify_agent.protocol import CreateExecutionBindingRequest, DestroyExecutionBindingRequest -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from clients.agent_backend.factory import create_agent_backend_client @@ -364,46 +364,68 @@ class AgentWorkspaceService: .order_by(AgentWorkspaceBinding.created_at) ).all() if not bindings: - logger.error( - "RETIRED Workspace has no Binding available for physical collection", - extra={"tenant_id": tenant_id, "workspace_id": workspace_id}, + raise AgentWorkspaceError( + f"RETIRED Workspace has no RETIRED Binding: tenant_id={tenant_id}, workspace_id={workspace_id}" ) - return anchor = bindings[0] - remaining_ids = [binding.id for binding in bindings[1:]] + remaining = [(binding.id, binding.backend_binding_ref) for binding in bindings[1:]] workspace_ref = workspace.backend_workspace_ref binding_ref = anchor.backend_binding_ref anchor_id = anchor.id + + failures: list[str] = [] + first_error: Exception | None = None with cls._client() as client: - client.destroy_execution_binding_sync( - DestroyExecutionBindingRequest( - binding_ref=binding_ref, - workspace_ref=workspace_ref, - destroy_workspace=True, + targets = [(anchor_id, binding_ref, workspace_ref, True)] + [ + (binding_id, backend_binding_ref, None, False) for binding_id, backend_binding_ref in remaining + ] + for binding_id, backend_binding_ref, target_workspace_ref, destroy_workspace in targets: + try: + client.destroy_execution_binding_sync( + DestroyExecutionBindingRequest( + binding_ref=backend_binding_ref, + workspace_ref=target_workspace_ref, + destroy_workspace=destroy_workspace, + ) + ) + except Exception as exc: + failures.append(binding_id) + if first_error is None: + first_error = exc + logger.exception( + "Failed to destroy retired Agent Workspace Binding", + extra={ + "tenant_id": tenant_id, + "workspace_id": workspace_id, + "binding_id": binding_id, + "destroy_workspace": destroy_workspace, + }, + ) + if failures: + if len(failures) == 1 and first_error is not None: + raise first_error + raise AgentWorkspaceError( + f"Failed to destroy {len(failures)} RETIRED Workspace Binding(s): {', '.join(failures)}" + ) from first_error + + binding_ids = [anchor_id, *(binding_id for binding_id, _binding_ref in remaining)] + with session_factory.create_session() as session: + session.execute( + delete(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.id.in_(binding_ids), + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.workspace_id == workspace_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, ) ) - with session_factory.create_session() as session: - stored_workspace = session.scalar( - select(AgentWorkspace).where( + session.execute( + delete(AgentWorkspace).where( AgentWorkspace.id == workspace_id, AgentWorkspace.tenant_id == tenant_id, AgentWorkspace.status == AgentWorkingResourceStatus.RETIRED, ) ) - stored_anchor = session.scalar( - select(AgentWorkspaceBinding).where( - AgentWorkspaceBinding.id == anchor_id, - AgentWorkspaceBinding.tenant_id == tenant_id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, - ) - ) - if stored_workspace is not None: - session.delete(stored_workspace) - if stored_anchor is not None: - session.delete(stored_anchor) session.commit() - for remaining_id in remaining_ids: - cls.collect_retired_binding(tenant_id=tenant_id, binding_id=remaining_id) @staticmethod def validate_binding_generation( diff --git a/api/services/agent_app_sandbox_service.py b/api/services/agent_app_sandbox_service.py index ad7b7b2fd85..f9b0149abf1 100644 --- a/api/services/agent_app_sandbox_service.py +++ b/api/services/agent_app_sandbox_service.py @@ -476,6 +476,7 @@ def _default_client_factory() -> Client: return create_agent_backend_client( base_url=base_url, api_token=dify_config.AGENT_BACKEND_API_TOKEN, + binding_file_download_timeout=dify_config.AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS, ) diff --git a/api/services/app_definition_query_service.py b/api/services/app_definition_query_service.py index 6ebd367a961..6815911b119 100644 --- a/api/services/app_definition_query_service.py +++ b/api/services/app_definition_query_service.py @@ -41,6 +41,7 @@ class AppSiteConfiguration(NamedTuple): input_placeholder: str | None custom_disclaimer: str | None default_language: str + prompt_public: bool show_workflow_steps: bool use_icon_as_answer_icon: bool diff --git a/api/services/app_dsl_service.py b/api/services/app_dsl_service.py index dd7f93818e4..a1b0d370489 100644 --- a/api/services/app_dsl_service.py +++ b/api/services/app_dsl_service.py @@ -19,7 +19,7 @@ from configs import dify_config from constants.dsl_version import CURRENT_APP_DSL_VERSION from core.file import remote_fetcher from core.plugin.entities.plugin import PluginDependency -from core.rbac import RBACPermission +from core.rbac import RBACPermission, RBACResourceScope from core.trigger.constants import ( TRIGGER_PLUGIN_NODE_TYPE, TRIGGER_SCHEDULE_NODE_TYPE, @@ -65,7 +65,6 @@ from services.errors.app import WorkflowNotFoundError from services.plugin.dependencies_analysis import DependenciesAnalysisService from services.workflow_draft_variable_service import WorkflowDraftVariableService from services.workflow_service import WorkflowService -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -228,9 +227,7 @@ class AppDslService: # If app_id is provided, check if it exists app = None if app_id: - stmt = select(App).where(App.id == app_id, App.tenant_id == account.current_tenant_id) - app = self._session.scalar(stmt) - + app = self._load_app_for_overwrite(account, app_id) if not app: return Import( id=import_id, @@ -368,8 +365,13 @@ class AppDslService: app = None if pending_data.app_id: - stmt = select(App).where(App.id == pending_data.app_id, App.tenant_id == account.current_tenant_id) - app = self._session.scalar(stmt) + app = self._load_app_for_overwrite(account, pending_data.app_id) + if not app: + return Import( + id=import_id, + status=ImportStatus.FAILED, + error="App not found", + ) # Create or update app app = self._create_or_update_app( @@ -430,6 +432,31 @@ class AppDslService: leaked_dependencies=leaked_dependencies, ) + def _load_app_for_overwrite(self, account: Account, app_id: str) -> App | None: + if account.current_tenant_id is None: + raise ValueError("Current tenant is not set") + if dify_config.RBAC_ENABLED and self._session.in_transaction(): + raise RuntimeError("App overwrite authorization requires a session without an active transaction") + rbac_allowed = not dify_config.RBAC_ENABLED or RBACService.CheckAccess.check( + account.current_tenant_id, + account.id, + scene=RBACPermission.APP_IMPORT_EXPORT_DSL, + resource_type=RBACResourceScope.APP, + resource_id=app_id, + ) + app = self._session.scalar( + select(App) + .where( + App.id == app_id, + App.tenant_id == account.current_tenant_id, + App.status == "normal", + ) + .execution_options(populate_existing=True) + ) + if app is not None and not rbac_allowed and app.maintainer != account.id: + raise NoPermissionError("You do not have permission to overwrite this app") + return app + @staticmethod def _ensure_agent_manage_permission(account: Account) -> None: """Importing an Agent DSL creates a roster Agent, which requires ``agent.manage``.""" @@ -601,16 +628,11 @@ class AppDslService: draft_workflow=draft_workflow, ) self._session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id=app.tenant_id, agent_ids=retirement_candidates, account_id=account.id, ) - enqueue_agent_resource_collection( - tenant_id=app.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) case AppMode.CHAT | AppMode.AGENT_CHAT | AppMode.COMPLETION: # Initialize model config model_config = data.get("model_config") diff --git a/api/services/app_service.py b/api/services/app_service.py index d7788ef6bf4..fdf356ac1a7 100644 --- a/api/services/app_service.py +++ b/api/services/app_service.py @@ -7,7 +7,7 @@ from typing import Any, Literal, NotRequired, TypedDict, cast, override import sqlalchemy as sa from pydantic import BaseModel, Field -from sqlalchemy import ColumnElement, select +from sqlalchemy import ColumnElement, delete, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session @@ -36,6 +36,8 @@ from models.agent import ( AgentStatus, AgentWorkingResourceStatus, AgentWorkspaceBinding, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, ) from models.model import App, AppMode, AppModelConfig, IconType, Site, load_annotation_reply_config from models.workflow import Workflow @@ -287,6 +289,13 @@ class AppService: ) -> App | None: return session.get(App, app_id) + @staticmethod + def get_normal_app_by_id( + app_id: str, + session: Session, + ) -> App | None: + return session.scalar(select(App).where(App.id == app_id, App.status == "normal").limit(1)) + @staticmethod def get_visible_app_by_id( app_id: str, @@ -741,7 +750,7 @@ class AppService: role: NotRequired[str | None] @staticmethod - def _get_backing_agent_for_update(app: App, *, session: Session) -> Agent | None: + def _get_backing_agent(app: App, *, session: Session) -> Agent | None: if app.mode != AppMode.AGENT: return None return session.scalar( @@ -784,7 +793,7 @@ class AppService: Role omission is intentional: ``role=None`` preserves the backing Agent's current role, while ``role=""`` explicitly clears it. """ - agent = self._get_backing_agent_for_update(app, session=session) + agent = self._get_backing_agent(app, session=session) if agent is None: return @@ -988,21 +997,49 @@ class AppService: return app def delete_app(self, app: App, *, session: Session) -> None: - """ - Delete app - :param app: App instance + """Delete an App and commit the passed session. + + The transaction releases all of a Workflow App's binding owners across + draft and published versions, archives a backing Roster Agent, retires + its resources, and deletes the App. Deleting a Roster Agent's backing + App does not remove bindings owned by external Workflows. + + After commit, the main App cleanup is published first, followed by + workflow-only Agent retirement and the Roster resource collector. Any + publication failure propagates. """ app_was_deleted.send(app) - backing_agent = self._get_backing_agent_for_update(app, session=session) - workflow_agent_ids = session.scalars( - select(Agent.id).where( - Agent.tenant_id == app.tenant_id, - Agent.app_id == app.id, - Agent.scope == AgentScope.WORKFLOW_ONLY, - Agent.status == AgentStatus.ACTIVE, + backing_agent = self._get_backing_agent(app, session=session) + workflow_agent_ids = set( + session.scalars( + select(Agent.id).where( + Agent.tenant_id == app.tenant_id, + Agent.app_id == app.id, + Agent.scope == AgentScope.WORKFLOW_ONLY, + Agent.status == AgentStatus.ACTIVE, + ) + ).all() + ) + if app.mode in (AppMode.WORKFLOW, AppMode.ADVANCED_CHAT): + workflow_agent_ids.update( + agent_id + for agent_id in session.scalars( + select(WorkflowAgentNodeBinding.agent_id).where( + WorkflowAgentNodeBinding.tenant_id == app.tenant_id, + WorkflowAgentNodeBinding.app_id == app.id, + WorkflowAgentNodeBinding.binding_type == WorkflowAgentBindingType.INLINE_AGENT, + WorkflowAgentNodeBinding.agent_id.is_not(None), + ) + ).all() + if agent_id + ) + session.execute( + delete(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.tenant_id == app.tenant_id, + WorkflowAgentNodeBinding.app_id == app.id, + ) ) - ).all() account_id = current_user.id if current_user else None if backing_agent is not None: now = naive_utc_now() @@ -1019,17 +1056,16 @@ class AppService: select(AgentWorkspaceBinding).where( AgentWorkspaceBinding.tenant_id == app.tenant_id, AgentWorkspaceBinding.agent_id == backing_agent.id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, ) ).all() for binding in bindings: - binding_id = AgentWorkspaceService.retire_binding( - session=session, - tenant_id=app.tenant_id, - binding_id=binding.id, - ) - if binding_id is not None: - retired_binding_ids.append(binding_id) + if binding.status == AgentWorkingResourceStatus.ACTIVE: + AgentWorkspaceService.retire_binding( + session=session, + tenant_id=app.tenant_id, + binding_id=binding.id, + ) + retired_binding_ids.append(binding.id) retired_snapshot_ids = AgentHomeSnapshotService.retire_all_for_agent( session=session, tenant_id=app.tenant_id, @@ -1044,7 +1080,16 @@ class AppService: session.delete(app) session.commit() - workflow_binding_ids, workflow_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + try: + remove_app_and_related_data_task.delay(tenant_id=app.tenant_id, app_id=app.id) + except Exception: + logger.exception( + "Failed to enqueue App cleanup", + extra={"tenant_id": app.tenant_id, "app_id": app.id}, + ) + raise + + WorkflowAgentRetirementService.retire_unowned( tenant_id=app.tenant_id, agent_ids=workflow_agent_ids, account_id=account_id, @@ -1052,8 +1097,9 @@ class AppService: enqueue_agent_resource_collection( tenant_id=app.tenant_id, workspace_ids=retired_workspace_ids, - binding_ids=[*retired_binding_ids, *workflow_binding_ids], - home_snapshot_ids=[*retired_snapshot_ids, *workflow_snapshot_ids], + binding_ids=retired_binding_ids, + home_snapshot_ids=retired_snapshot_ids, + purge_agent_ids=[backing_agent.id] if backing_agent is not None else [], ) # clean up web app settings @@ -1063,9 +1109,6 @@ class AppService: if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: BillingService.clean_billing_info_cache(app.tenant_id) - # Trigger asynchronous deletion of app and related data - remove_app_and_related_data_task.delay(tenant_id=app.tenant_id, app_id=app.id) - @staticmethod def get_app_code_by_id(app_id: str, *, session: Session) -> str: """ diff --git a/api/services/billing_service.py b/api/services/billing_service.py index e4c38f8874c..2ceaf144498 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -28,6 +28,9 @@ _http_client: httpx.Client = get_pooled_http_client( ) +EmailFreezeType = Literal["freeze", "email_domain_suspended"] + + class SubscriptionPlan(TypedDict): """Tenant subscriptionplan information.""" @@ -479,13 +482,26 @@ class BillingService: return cls._send_request("DELETE", "/account", params=params) @classmethod - def is_email_in_freeze(cls, email: str) -> bool: + def get_email_freeze_type(cls, email: str) -> EmailFreezeType | None: params = {"email": email} try: response = cls._send_request("GET", "/account/in-freeze", params=params) - return bool(response.get("data", False)) + if not response.get("data", False): + return None + + freeze_type = response.get("freeze_type") or response.get("freezeType") + if freeze_type in ("freeze", "email_domain_suspended"): + return freeze_type + + # Keep compatibility with older billing services that only return + # the boolean `data` field. + return "freeze" except Exception: - return False + return None + + @classmethod + def is_email_in_freeze(cls, email: str) -> bool: + return cls.get_email_freeze_type(email) is not None @classmethod def update_account_deletion_feedback(cls, email: str, feedback: str): diff --git a/api/services/data_migration/import_service.py b/api/services/data_migration/import_service.py index 48ff50ac655..960d58d92d8 100644 --- a/api/services/data_migration/import_service.py +++ b/api/services/data_migration/import_service.py @@ -18,6 +18,7 @@ import yaml from sqlalchemy import or_ from sqlalchemy.orm import Session, sessionmaker +from configs import dify_config from core.entities.mcp_provider import IdentityMode, MCPAuthentication, MCPConfiguration from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration from extensions.ext_database import db @@ -26,7 +27,6 @@ from models import Account, ApiToken, Tenant, TenantAccountJoin, TenantAccountRo from models.enums import ApiTokenType from models.model import App from models.tools import ApiToolProvider, MCPToolProvider, WorkflowToolProvider -from services.agent.retirement_service import WorkflowAgentRetirementService from services.app_dsl_service import AppDslService from services.data_migration.dependency_discovery_service import DependencyDiscoveryService from services.data_migration.entities import ( @@ -48,7 +48,6 @@ from services.tools.api_tools_manage_service import ApiToolManageService from services.tools.mcp_tools_manage_service import MCPToolManageService from services.tools.workflow_tools_manage_service import WorkflowToolManageService from services.workflow_service import WorkflowService -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection @dataclass(frozen=True) @@ -326,11 +325,14 @@ class MigrationImportService: ) -> str: import_service = AppDslService(session) if existing_app is not None: + existing_app_id = existing_app.id + if dify_config.RBAC_ENABLED: + session.commit() import_result = import_service.import_app( account=account, import_mode="yaml-content", yaml_content=dsl_content, - app_id=existing_app.id, + app_id=existing_app_id, ) else: import_app_id = app_id if self._should_preserve_source_app_id(options) else None @@ -713,7 +715,7 @@ class MigrationImportService: raise MigrationDataError(f"Referenced workflow app was not found in target tenant: {app_id}") if account_in_session is None: raise MigrationDataError(f"Operator account not found: {account.id}") - workflow, retirement_candidates = workflow_service.publish_workflow( + workflow = workflow_service.publish_workflow( session=session, app_model=app_in_session, account=account_in_session, @@ -723,16 +725,6 @@ class MigrationImportService: app_in_session.workflow_id = workflow.id app_in_session.updated_by = account.id app_in_session.updated_at = naive_utc_now() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( - tenant_id=target.tenant_id, - agent_ids=retirement_candidates, - account_id=account.id, - ) - enqueue_agent_resource_collection( - tenant_id=target.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) def _import_mcp_tools( self, diff --git a/api/services/entities/account_entities.py b/api/services/entities/account_entities.py new file mode 100644 index 00000000000..c3244d097a4 --- /dev/null +++ b/api/services/entities/account_entities.py @@ -0,0 +1,42 @@ +"""Framework-neutral contracts for Console account use cases.""" + +from dataclasses import dataclass +from datetime import datetime + + +@dataclass(frozen=True, slots=True) +class AccountSnapshot: + id: str + name: str + email: str + avatar: str | None + is_password_set: bool + interface_language: str | None + interface_theme: str | None + timezone: str | None + last_login_at: datetime | None + last_login_ip: str | None + status: str + initialized_at: datetime | None + created_at: datetime + + +@dataclass(frozen=True, slots=True) +class AccountProfileChanges: + name: str | None = None + avatar: str | None = None + interface_language: str | None = None + interface_theme: str | None = None + timezone: str | None = None + + def has_changes(self) -> bool: + return any( + value is not None + for value in ( + self.name, + self.avatar, + self.interface_language, + self.interface_theme, + self.timezone, + ) + ) diff --git a/api/services/errors/account.py b/api/services/errors/account.py index be421b631e6..090ecc373a0 100644 --- a/api/services/errors/account.py +++ b/api/services/errors/account.py @@ -9,6 +9,11 @@ class AccountRegisterError(BaseServiceError): pass +class EmailDomainSuspendedError(AccountRegisterError): + def __init__(self, description: str = "This email domain has been suspended."): + super().__init__(description) + + class AccountLoginError(BaseServiceError): pass diff --git a/api/services/feature_service.py b/api/services/feature_service.py index b882f5e0254..754413cee36 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -163,6 +163,10 @@ class FeatureService: def is_webapp_auth_enabled() -> bool: return dify_config.DEPLOYMENT_EDITION == DeploymentEdition.ENTERPRISE + @staticmethod + def is_trial_app_enabled() -> bool: + return dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and dify_config.ENABLE_TRIAL_APP + @classmethod def _fulfill_system_params_from_env(cls, system_features: feature_entities.SystemFeatureModel): system_features.enable_email_code_login = dify_config.ENABLE_EMAIL_CODE_LOGIN diff --git a/api/services/file_service.py b/api/services/file_service.py index 3ddc81ab41f..4497639eb25 100644 --- a/api/services/file_service.py +++ b/api/services/file_service.py @@ -20,6 +20,7 @@ from constants import ( VIDEO_EXTENSIONS, ) from core.rag.extractor.extract_processor import ExtractProcessor +from enums import DeploymentEdition from extensions.ext_storage import storage from extensions.storage.storage_type import StorageType from graphon.file import helpers as file_helpers @@ -179,6 +180,13 @@ class FileService: content_type=content_type, ) + def get_icon_url(self, file_id: str, tenant_id: str) -> str: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and ( + StorageType(dify_config.STORAGE_TYPE) == StorageType.S3 + ): + return self.get_file_presigned_url(file_id=file_id, tenant_id=tenant_id) + return file_helpers.get_signed_file_url(upload_file_id=file_id) + def upload_text(self, text: str, text_name: str, user_id: str, tenant_id: str) -> UploadFile: if len(text_name) > 200: text_name = text_name[:200] diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index 452e958ba6c..00daaa2bdd0 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -301,7 +301,11 @@ class RagPipelineService: return workflow def get_published_workflow_by_id(self, pipeline: Pipeline, workflow_id: str) -> Workflow | None: - """Fetch a published workflow snapshot by ID for restore operations.""" + """Fetch and lock a published Workflow snapshot for restoration. + + The source lock is held until the service transaction ends, preventing + concurrent deletion while restore copies its Workflow snapshot fields. + """ workflow = self._session.scalar( select(Workflow) .where( @@ -310,6 +314,7 @@ class RagPipelineService: Workflow.id == workflow_id, ) .limit(1) + .with_for_update() ) if workflow and workflow.version == Workflow.VERSION_DRAFT: raise IsDraftWorkflowError("source workflow must be published") @@ -419,7 +424,8 @@ class RagPipelineService: Pipelines reuse the shared draft-restore field copy helper, but still own the pipeline-specific flush/link step that wires a newly created draft - back onto ``pipeline.workflow_id``. + back onto ``pipeline.workflow_id``. The source version remains locked + through snapshot-field copy and commit. """ source_workflow = self.get_published_workflow_by_id(pipeline=pipeline, workflow_id=workflow_id) if not source_workflow: diff --git a/api/services/recommend_app/__init__.py b/api/services/recommend_app/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/api/services/recommend_app/buildin/__init__.py b/api/services/recommend_app/buildin/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/api/services/recommend_app/buildin/buildin_retrieval.py b/api/services/recommend_app/buildin/buildin_retrieval.py deleted file mode 100644 index d29d754b67e..00000000000 --- a/api/services/recommend_app/buildin/buildin_retrieval.py +++ /dev/null @@ -1,76 +0,0 @@ -import json -from os import path -from pathlib import Path -from typing import Any, override - -from flask import current_app -from sqlalchemy.orm import Session - -from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval -from services.recommend_app.recommend_app_base import RecommendAppRetrievalBase -from services.recommend_app.recommend_app_type import RecommendAppType - - -class BuildInRecommendAppRetrieval(RecommendAppRetrievalBase): - """ - Retrieval recommended app from buildin, the location is constants/recommended_apps.json - """ - - builtin_data: dict[str, Any] | None = None - - @override - def get_type(self) -> str: - return RecommendAppType.BUILDIN - - @override - def get_recommended_apps_and_categories(self, language: str, *, session: Session): - del session - result = self.fetch_recommended_apps_from_builtin(language) - return result - - @override - def get_learn_dify_apps(self, language: str, *, session: Session): - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db(language, session=session) - return result - - @override - def get_recommend_app_detail(self, app_id: str, *, session: Session): - del session - result = self.fetch_recommended_app_detail_from_builtin(app_id) - return result - - @classmethod - def _get_builtin_data(cls): - """ - Get builtin data. - :return: - """ - if cls.builtin_data: - return cls.builtin_data - - root_path = current_app.root_path - cls.builtin_data = json.loads( - Path(path.join(root_path, "constants", "recommended_apps.json")).read_text(encoding="utf-8") - ) - - return cls.builtin_data or {} - - @classmethod - def fetch_recommended_apps_from_builtin(cls, language: str): - """ - Fetch recommended apps from builtin. - :param language: language - :return: - """ - builtin_data: dict[str, dict[str, dict]] = cls._get_builtin_data() - return builtin_data.get("recommended_apps", {}).get(language, {}) - - @classmethod - def fetch_recommended_app_detail_from_builtin(cls, app_id: str) -> dict[str, Any] | None: - """ - Fetch recommended app detail from builtin. - :param app_id: App ID - :return: - """ - builtin_data: dict[str, dict[str, dict]] = cls._get_builtin_data() - return builtin_data.get("app_details", {}).get(app_id) diff --git a/api/services/recommend_app/category_order.py b/api/services/recommend_app/category_order.py deleted file mode 100644 index be6b112aa40..00000000000 --- a/api/services/recommend_app/category_order.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Apply Redis-backed category ordering for DB-backed Explore apps.""" - -import json -import logging -from collections.abc import Collection -from typing import Any - -from extensions.ext_redis import redis_client - -logger = logging.getLogger(__name__) - -EXPLORE_APP_CATEGORY_ORDER_KEY_PREFIX = "explore:apps:category_order" - - -def _category_order_key(language: str) -> str: - return f"{EXPLORE_APP_CATEGORY_ORDER_KEY_PREFIX}:{language}" - - -def get_explore_app_category_order(language: str) -> list[str]: - try: - raw_categories = redis_client.get(_category_order_key(language)) - except Exception: - logger.exception("Failed to read explore app category order from Redis.") - return [] - - if not raw_categories: - return [] - - if isinstance(raw_categories, bytes): - raw_categories = raw_categories.decode("utf-8") - - try: - categories: Any = json.loads(raw_categories) - except (TypeError, json.JSONDecodeError): - logger.warning("Invalid explore app category order payload for language %s.", language) - return [] - - if not isinstance(categories, list): - return [] - - return [category for category in categories if isinstance(category, str)] - - -def order_categories(categories: Collection[str], language: str) -> list[str]: - configured_order = get_explore_app_category_order(language) - if configured_order: - return configured_order - - return sorted(categories) diff --git a/api/services/recommend_app/database/__init__.py b/api/services/recommend_app/database/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/api/services/recommend_app/database/database_retrieval.py b/api/services/recommend_app/database/database_retrieval.py deleted file mode 100644 index 08d902fdeb5..00000000000 --- a/api/services/recommend_app/database/database_retrieval.py +++ /dev/null @@ -1,175 +0,0 @@ -from typing import Any, NotRequired, TypedDict, override - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from constants.languages import languages -from models.model import App, RecommendedApp -from services.app_dsl_service import AppDslService -from services.recommend_app.category_order import order_categories -from services.recommend_app.recommend_app_base import RecommendAppRetrievalBase -from services.recommend_app.recommend_app_type import RecommendAppType - - -class RecommendedAppItemDict(TypedDict): - id: str - app: App | None - app_id: str - description: Any - copyright: Any - privacy_policy: Any - custom_disclaimer: str - categories: list[str] - position: int - is_listed: bool - can_trial: NotRequired[bool] - - -class RecommendedAppsResultDict(TypedDict): - recommended_apps: list[RecommendedAppItemDict] - categories: list[str] - - -class RecommendedAppDetailDict(TypedDict): - id: str - name: str - icon: Any - icon_background: str | None - mode: str - export_data: str - - -class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase): - """ - Retrieval recommended app from database - """ - - @override - def get_recommended_apps_and_categories(self, language: str, *, session: Session) -> RecommendedAppsResultDict: - result = self.fetch_recommended_apps_from_db(language, session=session) - return result - - @override - def get_learn_dify_apps(self, language: str, *, session: Session) -> RecommendedAppsResultDict: - result = self.fetch_learn_dify_apps_from_db(language, session=session) - return result - - @override - def get_recommend_app_detail(self, app_id: str, *, session: Session) -> RecommendedAppDetailDict | None: - result = self.fetch_recommended_app_detail_from_db(app_id, session=session) - return result - - @override - def get_type(self) -> str: - return RecommendAppType.DATABASE - - @classmethod - def fetch_recommended_apps_from_db(cls, language: str, *, session: Session) -> RecommendedAppsResultDict: - """ - Fetch recommended apps from db. - :param language: language - :return: - """ - recommended_apps = cls._fetch_listed_recommended_apps(language, session=session) - - if len(recommended_apps) == 0: - recommended_apps = cls._fetch_listed_recommended_apps(languages[0], session=session) - - return cls._format_recommended_apps(recommended_apps, language) - - @classmethod - def fetch_learn_dify_apps_from_db(cls, language: str, *, session: Session) -> RecommendedAppsResultDict: - """ - Fetch listed recommended apps explicitly marked for the Learn Dify section. - :param language: language - :return: - """ - recommended_apps = cls._fetch_listed_recommended_apps(language, session=session, is_learn_dify=True) - - if len(recommended_apps) == 0 and language != languages[0]: - recommended_apps = cls._fetch_listed_recommended_apps(languages[0], session=session, is_learn_dify=True) - - return cls._format_recommended_apps(recommended_apps, language) - - @classmethod - def _fetch_listed_recommended_apps( - cls, language: str, *, session: Session, is_learn_dify: bool | None = None - ) -> list[RecommendedApp]: - filters = [RecommendedApp.is_listed.is_(True), RecommendedApp.language == language] - if is_learn_dify is not None: - filters.append(RecommendedApp.is_learn_dify.is_(is_learn_dify)) - - return list(session.scalars(select(RecommendedApp).where(*filters)).all()) - - @classmethod - def _format_recommended_apps( - cls, recommended_apps: list[RecommendedApp], language: str - ) -> RecommendedAppsResultDict: - """ - Serialize DB recommended app rows into the Explore list response shape. - :param recommended_apps: recommended app rows - :param language: language used for category ordering - :return: - """ - - categories = set() - recommended_apps_result: list[RecommendedAppItemDict] = [] - for recommended_app in recommended_apps: - app = recommended_app.app - if not app or not app.is_public: - continue - - site = app.site - if not site: - continue - - app_categories = recommended_app.categories or [] - recommended_app_result: RecommendedAppItemDict = { - "id": recommended_app.id, - "app": recommended_app.app, - "app_id": recommended_app.app_id, - "description": site.description, - "copyright": site.copyright, - "privacy_policy": site.privacy_policy, - "custom_disclaimer": site.custom_disclaimer, - "categories": app_categories, - "position": recommended_app.position, - "is_listed": recommended_app.is_listed, - } - recommended_apps_result.append(recommended_app_result) - - categories.update(app_categories) - - return RecommendedAppsResultDict( - recommended_apps=recommended_apps_result, - categories=order_categories(categories, language), - ) - - @classmethod - def fetch_recommended_app_detail_from_db(cls, app_id: str, *, session: Session) -> RecommendedAppDetailDict | None: - """ - Fetch recommended app detail from db. - :param app_id: App ID - :return: - """ - # is in public recommended list - recommended_app = session.scalar( - select(RecommendedApp).where(RecommendedApp.is_listed == True, RecommendedApp.app_id == app_id).limit(1) - ) - - if not recommended_app: - return None - - # get app detail - app_model = session.get(App, app_id) - if not app_model or not app_model.is_public: - return None - - return RecommendedAppDetailDict( - id=app_model.id, - name=app_model.name, - icon=app_model.icon, - icon_background=app_model.icon_background, - mode=app_model.mode, - export_data=AppDslService.export_dsl(app_model=app_model, session=session), - ) diff --git a/api/services/recommend_app/recommend_app_base.py b/api/services/recommend_app/recommend_app_base.py deleted file mode 100644 index 821ad476c42..00000000000 --- a/api/services/recommend_app/recommend_app_base.py +++ /dev/null @@ -1,15 +0,0 @@ -from typing import Any, Protocol - -from sqlalchemy.orm import Session - - -class RecommendAppRetrievalBase(Protocol): - """Interface for recommend app retrieval.""" - - def get_recommended_apps_and_categories(self, language: str, *, session: Session) -> Any: ... - - def get_learn_dify_apps(self, language: str, *, session: Session) -> Any: ... - - def get_recommend_app_detail(self, app_id: str, *, session: Session) -> Any: ... - - def get_type(self) -> str: ... diff --git a/api/services/recommend_app/recommend_app_factory.py b/api/services/recommend_app/recommend_app_factory.py deleted file mode 100644 index e53667c0b06..00000000000 --- a/api/services/recommend_app/recommend_app_factory.py +++ /dev/null @@ -1,23 +0,0 @@ -from services.recommend_app.buildin.buildin_retrieval import BuildInRecommendAppRetrieval -from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval -from services.recommend_app.recommend_app_base import RecommendAppRetrievalBase -from services.recommend_app.recommend_app_type import RecommendAppType -from services.recommend_app.remote.remote_retrieval import RemoteRecommendAppRetrieval - - -class RecommendAppRetrievalFactory: - @staticmethod - def get_recommend_app_factory(mode: str) -> type[RecommendAppRetrievalBase]: - match mode: - case RecommendAppType.REMOTE: - return RemoteRecommendAppRetrieval - case RecommendAppType.DATABASE: - return DatabaseRecommendAppRetrieval - case RecommendAppType.BUILDIN: - return BuildInRecommendAppRetrieval - case _: - raise ValueError(f"invalid fetch recommended apps mode: {mode}") - - @staticmethod - def get_buildin_recommend_app_retrieval(): - return BuildInRecommendAppRetrieval diff --git a/api/services/recommend_app/recommend_app_type.py b/api/services/recommend_app/recommend_app_type.py deleted file mode 100644 index e60e435b3a0..00000000000 --- a/api/services/recommend_app/recommend_app_type.py +++ /dev/null @@ -1,7 +0,0 @@ -from enum import StrEnum - - -class RecommendAppType(StrEnum): - REMOTE = "remote" - BUILDIN = "builtin" - DATABASE = "db" diff --git a/api/services/recommend_app/remote/__init__.py b/api/services/recommend_app/remote/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/api/services/recommend_app/remote/remote_retrieval.py b/api/services/recommend_app/remote/remote_retrieval.py deleted file mode 100644 index d30306382ff..00000000000 --- a/api/services/recommend_app/remote/remote_retrieval.py +++ /dev/null @@ -1,173 +0,0 @@ -import logging -import threading -from typing import Any, override - -import httpx -from cachetools import TTLCache -from flask import has_request_context, request -from sqlalchemy.orm import Session - -from configs import dify_config -from services.recommend_app.buildin.buildin_retrieval import BuildInRecommendAppRetrieval -from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval -from services.recommend_app.recommend_app_base import RecommendAppRetrievalBase -from services.recommend_app.recommend_app_type import RecommendAppType - -logger = logging.getLogger(__name__) - -_REMOTE_FETCH_CACHE_MAXSIZE = 64 -_remote_fetch_cache: TTLCache[tuple[str, str], dict[str, Any]] | None = None -_remote_fetch_cache_ttl: int | None = None -_remote_fetch_cache_lock = threading.Lock() - - -def _current_origin_headers() -> dict[str, str]: - origin = request.headers.get("Origin") if has_request_context() else None - if origin: - return {"Origin": origin} - - console_web_url = getattr(dify_config, "CONSOLE_WEB_URL", "") - if not isinstance(console_web_url, str) or not console_web_url: - return {} - return {"Origin": console_web_url} - - -def _remote_fetch_cache_key(url: str, headers: dict[str, str]) -> tuple[str, str]: - return url, headers.get("Origin", "") - - -def _hosted_fetch_cache_ttl() -> int: - ttl = dify_config.HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL - if isinstance(ttl, int) and not isinstance(ttl, bool): - return ttl - return 600 - - -def _get_remote_fetch_cache() -> TTLCache[tuple[str, str], dict[str, Any]] | None: - ttl = _hosted_fetch_cache_ttl() - if ttl <= 0: - return None - - global _remote_fetch_cache, _remote_fetch_cache_ttl - if _remote_fetch_cache is None or _remote_fetch_cache_ttl != ttl: - with _remote_fetch_cache_lock: - if _remote_fetch_cache is None or _remote_fetch_cache_ttl != ttl: - _remote_fetch_cache = TTLCache(maxsize=_REMOTE_FETCH_CACHE_MAXSIZE, ttl=ttl) - _remote_fetch_cache_ttl = ttl - return _remote_fetch_cache - - -def clear_remote_fetch_cache() -> None: - """Reset the in-memory remote fetch cache (used by tests).""" - global _remote_fetch_cache, _remote_fetch_cache_ttl - with _remote_fetch_cache_lock: - _remote_fetch_cache = None - _remote_fetch_cache_ttl = None - - -def _fetch_remote_payload(url: str) -> tuple[int, dict[str, Any] | None]: - headers = _current_origin_headers() - cache_key = _remote_fetch_cache_key(url, headers) - cache = _get_remote_fetch_cache() - if cache is not None: - with _remote_fetch_cache_lock: - cached = cache.get(cache_key) - if cached is not None: - return 200, cached - - response = httpx.get(url, headers=headers, timeout=httpx.Timeout(10.0, connect=3.0)) - status_code = response.status_code - if status_code != 200: - return status_code, None - - result: dict[str, Any] = response.json() - if cache is not None: - with _remote_fetch_cache_lock: - cache[cache_key] = result - return status_code, result - - -class RemoteRecommendAppRetrieval(RecommendAppRetrievalBase): - """ - Retrieval recommended app from dify official. - - The remote `/apps` payload is already curated for display, including category order. - Keep the response order intact so Explore matches the template service. - """ - - @override - def get_recommend_app_detail(self, app_id: str, *, session: Session): - del session - try: - result = self.fetch_recommended_app_detail_from_dify_official(app_id) - except Exception as e: - logger.warning("fetch recommended app detail from dify official failed: %s, switch to built-in.", e) - result = BuildInRecommendAppRetrieval.fetch_recommended_app_detail_from_builtin(app_id) - return result - - @override - def get_recommended_apps_and_categories(self, language: str, *, session: Session): - del session - try: - result = self.fetch_recommended_apps_from_dify_official(language) - except Exception as e: - logger.warning("fetch recommended apps from dify official failed: %s, switch to built-in.", e) - result = BuildInRecommendAppRetrieval.fetch_recommended_apps_from_builtin(language) - return result - - @override - def get_learn_dify_apps(self, language: str, *, session: Session): - try: - result = self.fetch_learn_dify_apps_from_dify_official(language) - except Exception as e: - logger.warning("fetch learn dify apps from dify official failed: %s, switch to database.", e) - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db(language, session=session) - return result - - @override - def get_type(self) -> str: - return RecommendAppType.REMOTE - - @classmethod - def fetch_recommended_app_detail_from_dify_official(cls, app_id: str) -> dict[str, Any] | None: - """ - Fetch recommended app detail from dify official. - :param app_id: App ID - :return: - """ - domain = dify_config.HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN - url = f"{domain}/apps/{app_id}" - status_code, data = _fetch_remote_payload(url) - if status_code != 200: - return None - return data - - @classmethod - def fetch_recommended_apps_from_dify_official(cls, language: str): - """ - Fetch recommended apps from dify official. - :param language: language - :return: - """ - domain = dify_config.HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN - url = f"{domain}/apps?language={language}" - status_code, result = _fetch_remote_payload(url) - if status_code != 200: - raise ValueError(f"fetch recommended apps failed, status code: {status_code}") - - return result - - @classmethod - def fetch_learn_dify_apps_from_dify_official(cls, language: str): - """ - Fetch Learn Dify apps from dify official. - :param language: language - :return: - """ - domain = dify_config.HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN - url = f"{domain}/apps/learn-dify?language={language}" - status_code, result = _fetch_remote_payload(url) - if status_code != 200: - raise ValueError(f"fetch learn dify apps failed, status code: {status_code}") - - return result diff --git a/api/services/recommended_app_catalog_gateway.py b/api/services/recommended_app_catalog_gateway.py new file mode 100644 index 00000000000..798ec9756a5 --- /dev/null +++ b/api/services/recommended_app_catalog_gateway.py @@ -0,0 +1,376 @@ +"""Typed remote and built-in adapters for the recommended app catalog.""" + +import json +import logging +import threading +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import cast, override + +import httpx +from cachetools import TTLCache + +from configs import dify_config +from services.recommended_app_query_service import ( + RecommendedAppCatalogPage, + RecommendedAppCatalogQuery, + RecommendedAppDetailRecord, + RecommendedAppInfoRecord, + RecommendedAppRecord, +) + +logger = logging.getLogger(__name__) + +_BUILTIN_FALLBACK_LANGUAGE = "en-US" +_BUILTIN_CATALOG_PATH = Path(__file__).resolve().parents[1] / "constants" / "recommended_apps.json" +_REMOTE_FETCH_CACHE_MAXSIZE = 64 +_remote_fetch_cache: TTLCache[tuple[str, str], object] | None = None +_remote_fetch_cache_ttl: int | None = None +_remote_fetch_cache_lock = threading.Lock() + + +def _hosted_fetch_cache_ttl() -> int: + ttl = dify_config.HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL + if isinstance(ttl, int) and not isinstance(ttl, bool): + return ttl + return 600 + + +def _get_remote_fetch_cache() -> TTLCache[tuple[str, str], object] | None: + ttl = _hosted_fetch_cache_ttl() + if ttl <= 0: + return None + + global _remote_fetch_cache, _remote_fetch_cache_ttl + if _remote_fetch_cache is None or _remote_fetch_cache_ttl != ttl: + with _remote_fetch_cache_lock: + if _remote_fetch_cache is None or _remote_fetch_cache_ttl != ttl: + _remote_fetch_cache = TTLCache(maxsize=_REMOTE_FETCH_CACHE_MAXSIZE, ttl=ttl) + _remote_fetch_cache_ttl = ttl + return _remote_fetch_cache + + +def clear_remote_fetch_cache() -> None: + """Reset the in-memory remote fetch cache (used by tests).""" + global _remote_fetch_cache, _remote_fetch_cache_ttl + with _remote_fetch_cache_lock: + _remote_fetch_cache = None + _remote_fetch_cache_ttl = None + + +class _RecommendedAppSourceUnavailableError(Exception): + pass + + +class BuiltinRecommendedAppCatalogGateway(RecommendedAppCatalogQuery): + def __init__(self) -> None: + self._data: Mapping[str, object] | None = None + + @override + def list_recommended(self, language: str) -> RecommendedAppCatalogPage: + return _map_recommended_page(self._raw_page(language)) + + @override + def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: + return _map_learn_dify_page(self._raw_learn_dify_page(language)) + + @override + def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: + detail = self._raw_detail(app_id) + if detail is None: + return None + return _map_detail(_as_mapping(detail, field="recommended app detail")) + + @override + def contains(self, app_id: str) -> bool: + return self._raw_detail(app_id) is not None + + def _raw_page(self, language: str) -> Mapping[str, object]: + pages = _as_mapping(self._get_data().get("recommended_apps", {}), field="recommended_apps") + return _as_mapping(pages.get(language, {}), field="recommended app page") + + def _raw_learn_dify_page(self, language: str) -> Mapping[str, object]: + apps = self._raw_learn_dify_apps(language) + if not apps and language != _BUILTIN_FALLBACK_LANGUAGE: + apps = self._raw_learn_dify_apps(_BUILTIN_FALLBACK_LANGUAGE) + return {"recommended_apps": apps} + + def _raw_learn_dify_apps(self, language: str) -> tuple[object, ...]: + page = self._raw_page(language) + return tuple( + app + for app in _as_sequence(page.get("recommended_apps", ()), field="apps") + if _as_mapping(app, field="recommended app").get("is_learn_dify") is True + ) + + def _raw_detail(self, app_id: str) -> object | None: + details = _as_mapping(self._get_data().get("app_details", {}), field="app_details") + return details.get(app_id) + + def _get_data(self) -> Mapping[str, object]: + if self._data is None: + loaded = json.loads(_BUILTIN_CATALOG_PATH.read_text(encoding="utf-8")) + self._data = _as_mapping(loaded, field="built-in recommended app catalog") + return self._data + + +class RemoteRecommendedAppCatalogGateway(RecommendedAppCatalogQuery): + @override + def list_recommended(self, language: str) -> RecommendedAppCatalogPage: + result = self._fetch(lambda: self._fetch_page(language)) + return _map_recommended_page(_as_mapping(result, field="recommended app page")) + + @override + def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: + result = self._fetch(lambda: self._fetch_learn_dify_page(language)) + return _map_learn_dify_page(_as_mapping(result, field="Learn Dify app page")) + + @override + def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: + detail = self._fetch(lambda: self._fetch_detail(app_id)) + if detail is None: + return None + return _map_detail(_as_mapping(detail, field="recommended app detail")) + + @override + def contains(self, app_id: str) -> bool: + detail = self._fetch(lambda: self._fetch_detail(app_id)) + return detail is not None + + def _fetch_detail(self, app_id: str) -> object | None: + status_code, detail = self._get_payload(f"/apps/{app_id}") + if status_code != 200: + # Preserve the legacy detail contract: only request or decoding + # failures use the bundled fallback; HTTP responses are authoritative. + return None + return detail + + def _fetch_page(self, language: str) -> object: + status_code, page = self._get_payload(f"/apps?language={language}") + if status_code != 200: + raise ValueError(f"fetch recommended apps failed, status code: {status_code}") + return page + + def _fetch_learn_dify_page(self, language: str) -> object: + status_code, page = self._get_payload(f"/apps/learn-dify?language={language}") + if status_code != 200: + raise ValueError(f"fetch learn dify apps failed, status code: {status_code}") + return page + + @staticmethod + def _get_payload(path: str) -> tuple[int, object]: + origin = dify_config.CONSOLE_WEB_URL + + url = f"{dify_config.HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN}{path}" + headers = {"Origin": origin} if origin else {} + cache_key = (url, origin) + cache = _get_remote_fetch_cache() + if cache is not None: + with _remote_fetch_cache_lock: + cached = cache.get(cache_key) + if cached is not None: + return 200, cached + + response = httpx.get( + url, + headers=headers, + timeout=httpx.Timeout(10.0, connect=3.0), + ) + if response.status_code != 200: + return response.status_code, None + + result = response.json() + if cache is not None: + with _remote_fetch_cache_lock: + cache[cache_key] = result + return response.status_code, result + + @staticmethod + def _fetch[T](fetch: Callable[[], T]) -> T: + try: + return fetch() + except Exception as error: + raise _RecommendedAppSourceUnavailableError(str(error)) from error + + +class RecommendedAppCatalogRouter(RecommendedAppCatalogQuery): + def __init__( + self, + *, + remote: RecommendedAppCatalogQuery, + database: RecommendedAppCatalogQuery, + builtin: RecommendedAppCatalogQuery, + ) -> None: + self._remote = remote + self._builtin = builtin + self._sources: dict[str, RecommendedAppCatalogQuery] = { + "remote": remote, + "db": database, + "builtin": builtin, + } + + @override + def list_recommended(self, language: str) -> RecommendedAppCatalogPage: + source = self._source() + if source is not self._remote: + page = source.list_recommended(language) + else: + try: + page = self._remote.list_recommended(language) + except _RecommendedAppSourceUnavailableError as error: + logger.warning("fetch recommended apps from dify official failed: %s, switch to built-in.", error) + page = self._builtin.list_recommended(language) + + if not page.recommended_apps: + return self._builtin.list_recommended(_BUILTIN_FALLBACK_LANGUAGE) + return page + + @override + def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: + source = self._source() + if source is not self._remote: + return source.list_learn_dify(language) + try: + return self._remote.list_learn_dify(language) + except _RecommendedAppSourceUnavailableError as error: + logger.warning("fetch learn dify apps from dify official failed: %s, switch to built-in.", error) + return self._builtin.list_learn_dify(language) + + @override + def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: + source = self._source() + if source is not self._remote: + return source.get_detail(app_id) + try: + return self._remote.get_detail(app_id) + except _RecommendedAppSourceUnavailableError as error: + logger.warning("fetch recommended app detail from dify official failed: %s, switch to built-in.", error) + return self._builtin.get_detail(app_id) + + @override + def contains(self, app_id: str) -> bool: + source = self._source() + if source is not self._remote: + return source.contains(app_id) + try: + return self._remote.contains(app_id) + except _RecommendedAppSourceUnavailableError as error: + logger.warning("fetch recommended app detail from dify official failed: %s, switch to built-in.", error) + return self._builtin.contains(app_id) + + def _source(self) -> RecommendedAppCatalogQuery: + mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE + try: + return self._sources[mode] + except KeyError: + raise ValueError(f"invalid fetch recommended apps mode: {mode}") from None + + +def _map_recommended_page(source: Mapping[str, object]) -> RecommendedAppCatalogPage: + if not source.get("recommended_apps"): + return RecommendedAppCatalogPage(recommended_apps=(), categories=()) + return _map_page(source) + + +def _map_page(source: Mapping[str, object]) -> RecommendedAppCatalogPage: + return RecommendedAppCatalogPage( + recommended_apps=tuple(_map_app(app) for app in _as_sequence(source["recommended_apps"], field="apps")), + categories=_as_string_tuple(source["categories"], field="categories"), + ) + + +def _map_learn_dify_page(source: Mapping[str, object]) -> RecommendedAppCatalogPage: + return RecommendedAppCatalogPage( + recommended_apps=tuple(_map_app(app) for app in _as_sequence(source["recommended_apps"], field="apps")), + categories=(), + ) + + +def _map_app(source: object) -> RecommendedAppRecord: + source = _as_mapping(source, field="recommended app") + app_id = source["app_id"] + if not isinstance(app_id, str): + raise TypeError("app_id must be a string") + + app_source = source.get("app") + if app_source is not None: + app_source = _as_mapping(app_source, field="app") + + return RecommendedAppRecord( + app=_map_app_info(app_source), + app_id=app_id, + description=cast(str | None, source.get("description")), + copyright=cast(str | None, source.get("copyright")), + privacy_policy=cast(str | None, source.get("privacy_policy")), + custom_disclaimer=cast(str | None, source.get("custom_disclaimer")), + categories=_as_string_tuple(source.get("categories", ()), field="categories"), + position=cast(int | None, source.get("position")), + is_listed=cast(bool | None, source.get("is_listed")), + ) + + +def _map_app_info(source: Mapping[str, object] | None) -> RecommendedAppInfoRecord | None: + if source is None: + return None + app_id = source["id"] + if not isinstance(app_id, str): + raise TypeError("app.id must be a string") + return RecommendedAppInfoRecord( + id=app_id, + name=cast(str | None, source.get("name")), + mode=_enum_string(source.get("mode"), field="app.mode"), + icon=cast(str | None, source.get("icon")), + icon_type=_enum_string(source.get("icon_type"), field="app.icon_type"), + icon_background=cast(str | None, source.get("icon_background")), + ) + + +def _map_detail(source: Mapping[str, object]) -> RecommendedAppDetailRecord: + app_id = source["id"] + name = source["name"] + export_data = source["export_data"] + if not isinstance(app_id, str): + raise TypeError("id must be a string") + if not isinstance(name, str): + raise TypeError("name must be a string") + if not isinstance(export_data, str): + raise TypeError("export_data must be a string") + + mode = _enum_string(source["mode"], field="mode") + if mode is None: + raise TypeError("mode must be a string or string enum") + return RecommendedAppDetailRecord( + id=app_id, + name=name, + icon=cast(str | None, source.get("icon")), + icon_background=cast(str | None, source.get("icon_background")), + mode=mode, + export_data=export_data, + ) + + +def _as_mapping(value: object, *, field: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise TypeError(f"{field} must be a mapping") + return cast(Mapping[str, object], value) + + +def _as_sequence(value: object, *, field: str) -> Sequence[object]: + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError(f"{field} must be a sequence") + return cast(Sequence[object], value) + + +def _as_string_tuple(value: object, *, field: str) -> tuple[str, ...]: + values = _as_sequence(value, field=field) + if not all(isinstance(item, str) for item in values): + raise TypeError(f"{field} must contain only strings") + return cast(tuple[str, ...], tuple(values)) + + +def _enum_string(value: object, *, field: str) -> str | None: + if value is None: + return None + if isinstance(value, str): + return str(value) + raise TypeError(f"{field} must be a string or string enum") diff --git a/api/services/recommended_app_query_service.py b/api/services/recommended_app_query_service.py new file mode 100644 index 00000000000..5b79f40d505 --- /dev/null +++ b/api/services/recommended_app_query_service.py @@ -0,0 +1,186 @@ +"""Application service for querying the recommended app catalog.""" + +from collections.abc import Sequence, Set +from typing import NamedTuple, Protocol + +from constants.languages import languages + + +class RecommendedAppInfoRecord(NamedTuple): + id: str + name: str | None + mode: str | None + icon: str | None + icon_type: str | None + icon_background: str | None + + +class RecommendedAppRecord(NamedTuple): + app: RecommendedAppInfoRecord | None + app_id: str + description: str | None + copyright: str | None + privacy_policy: str | None + custom_disclaimer: str | None + categories: tuple[str, ...] + position: int | None + is_listed: bool | None + + +class RecommendedAppCatalogPage(NamedTuple): + recommended_apps: tuple[RecommendedAppRecord, ...] + categories: tuple[str, ...] + + +class RecommendedAppDetailRecord(NamedTuple): + id: str + name: str + icon: str | None + icon_background: str | None + mode: str + export_data: str + + +class RecommendedAppCatalogQuery(Protocol): + """Read from the recommended-app catalog.""" + + def list_recommended(self, language: str) -> RecommendedAppCatalogPage: ... + + def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: ... + + def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: ... + + def contains(self, app_id: str) -> bool: ... + + +class TrialAppQuery(Protocol): + def existing_ids(self, app_ids: Sequence[str]) -> Set[str]: ... + + +class RecommendedAppSummary(NamedTuple): + app: RecommendedAppInfoRecord | None + app_id: str + description: str | None + copyright: str | None + privacy_policy: str | None + custom_disclaimer: str | None + categories: tuple[str, ...] + position: int | None + is_listed: bool | None + can_trial: bool + + +class RecommendedAppListResult(NamedTuple): + recommended_apps: tuple[RecommendedAppSummary, ...] + categories: tuple[str, ...] + + +class LearnDifyAppListResult(NamedTuple): + recommended_apps: tuple[RecommendedAppSummary, ...] + + +class RecommendedAppDetailSummary(NamedTuple): + id: str + name: str + icon: str | None + icon_background: str | None + mode: str + export_data: str + can_trial: bool + + +class RecommendedAppNotFoundError(Exception): + pass + + +class RecommendedAppQueryService: + def __init__( + self, + *, + catalog: RecommendedAppCatalogQuery, + trial_apps: TrialAppQuery, + trial_enabled: bool, + ) -> None: + self._catalog = catalog + self._trial_apps = trial_apps + self._trial_enabled = trial_enabled + + def is_trial_enabled(self) -> bool: + return self._trial_enabled + + def is_previewable(self, app_id: str) -> bool: + if app_id in self._trial_apps.existing_ids((app_id,)): + return True + return self._catalog.contains(app_id) + + def list_recommended( + self, + *, + requested_language: str | None, + interface_language: str | None, + ) -> RecommendedAppListResult: + language = self._resolve_language(requested_language, interface_language) + page = self._catalog.list_recommended(language) + + return RecommendedAppListResult( + recommended_apps=self._with_trial_status(page.recommended_apps), + categories=page.categories, + ) + + def list_learn_dify( + self, + *, + requested_language: str | None, + interface_language: str | None, + ) -> LearnDifyAppListResult: + language = self._resolve_language(requested_language, interface_language) + page = self._catalog.list_learn_dify(language) + return LearnDifyAppListResult(recommended_apps=self._with_trial_status(page.recommended_apps)) + + def get_detail(self, app_id: str) -> RecommendedAppDetailSummary: + detail = self._catalog.get_detail(app_id) + if detail is None: + raise RecommendedAppNotFoundError + + can_trial = False + if self._trial_enabled: + can_trial = detail.id in self._trial_apps.existing_ids((detail.id,)) + + return RecommendedAppDetailSummary( + id=detail.id, + name=detail.name, + icon=detail.icon, + icon_background=detail.icon_background, + mode=detail.mode, + export_data=detail.export_data, + can_trial=can_trial, + ) + + def _with_trial_status(self, apps: Sequence[RecommendedAppRecord]) -> tuple[RecommendedAppSummary, ...]: + trial_app_ids: Set[str] = set() + if self._trial_enabled: + trial_app_ids = self._trial_apps.existing_ids([app.app_id for app in apps]) + + return tuple( + RecommendedAppSummary( + app=app.app, + app_id=app.app_id, + description=app.description, + copyright=app.copyright, + privacy_policy=app.privacy_policy, + custom_disclaimer=app.custom_disclaimer, + categories=app.categories, + position=app.position, + is_listed=app.is_listed, + can_trial=app.app_id in trial_app_ids, + ) + for app in apps + ) + + @staticmethod + def _resolve_language(requested_language: str | None, interface_language: str | None) -> str: + if requested_language and requested_language in languages: + return requested_language + if interface_language: + return interface_language + return languages[0] diff --git a/api/services/recommended_app_service.py b/api/services/recommended_app_service.py deleted file mode 100644 index 7b09e2b4005..00000000000 --- a/api/services/recommended_app_service.py +++ /dev/null @@ -1,119 +0,0 @@ -from typing import Any - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from configs import dify_config -from enums import DeploymentEdition -from models.model import AccountTrialAppRecord, App, TrialApp -from services.recommend_app.recommend_app_factory import RecommendAppRetrievalFactory - - -class RecommendedAppService: - """Own recommended app retrieval and Cloud-only trial eligibility.""" - - @staticmethod - def is_trial_app_enabled() -> bool: - """Return whether trial execution is enabled for this deployment.""" - return dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and dify_config.ENABLE_TRIAL_APP - - @classmethod - def get_app(cls, app_id: str, *, session: Session) -> App | None: - """Return a normal app only when it belongs to the recommended catalog.""" - mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE - retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)() - recommended_app_detail = retrieval_instance.get_recommend_app_detail(app_id, session=session) - if recommended_app_detail is None: - return None - - return session.scalar(select(App).where(App.id == app_id, App.status == "normal").limit(1)) - - @classmethod - def get_recommended_apps_and_categories(cls, language: str, *, session: Session): - """ - Get recommended apps and categories. - :param language: language - :return: - """ - mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE - retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)() - result = retrieval_instance.get_recommended_apps_and_categories(language, session=session) - if not result.get("recommended_apps"): - result = ( - RecommendAppRetrievalFactory.get_buildin_recommend_app_retrieval().fetch_recommended_apps_from_builtin( - "en-US" - ) - ) - - apps = result["recommended_apps"] - trial_app_ids = ( - cls._get_trial_app_ids(session, [app["app_id"] for app in apps]) if cls.is_trial_app_enabled() else set() - ) - for app in apps: - app["can_trial"] = app["app_id"] in trial_app_ids - return result - - @classmethod - def get_learn_dify_apps(cls, language: str, *, session: Session) -> dict[str, Any]: - """ - Get recommended apps marked for the Learn Dify section. - :param language: language - :return: - """ - mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE - retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)() - result = retrieval_instance.get_learn_dify_apps(language, session=session) - - apps = result["recommended_apps"] - trial_app_ids = ( - cls._get_trial_app_ids(session, [app["app_id"] for app in apps]) if cls.is_trial_app_enabled() else set() - ) - for app in apps: - app["can_trial"] = app["app_id"] in trial_app_ids - - return {"recommended_apps": apps} - - @classmethod - def get_recommend_app_detail(cls, app_id: str, *, session: Session) -> dict[str, Any] | None: - """ - Get recommend app detail. - :param app_id: app id - :return: - """ - mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE - retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)() - result: dict[str, Any] | None = retrieval_instance.get_recommend_app_detail(app_id, session=session) - if result is None: - return None - result["can_trial"] = cls.is_trial_app_enabled() and cls._can_trial_app(session, result["id"]) - return result - - @classmethod - def add_trial_app_record(cls, app_id: str, account_id: str, *, session: Session): - """ - Add trial app record. - :param app_id: app id - :return: - """ - account_trial_app_record = session.scalar( - select(AccountTrialAppRecord) - .where(AccountTrialAppRecord.app_id == app_id, AccountTrialAppRecord.account_id == account_id) - .limit(1) - ) - if account_trial_app_record: - account_trial_app_record.count += 1 - session.commit() - else: - session.add(AccountTrialAppRecord(app_id=app_id, count=1, account_id=account_id)) - session.commit() - - @staticmethod - def _can_trial_app(session: Session, app_id: str) -> bool: - trial_app_model = session.scalar(select(TrialApp).where(TrialApp.app_id == app_id).limit(1)) - return trial_app_model is not None - - @staticmethod - def _get_trial_app_ids(session: Session, app_ids: list[str]) -> set[str]: - if not app_ids: - return set() - return set(session.scalars(select(TrialApp.app_id).where(TrialApp.app_id.in_(app_ids))).all()) diff --git a/api/services/snippet_dsl_service.py b/api/services/snippet_dsl_service.py index 22f495a2370..78361700ad4 100644 --- a/api/services/snippet_dsl_service.py +++ b/api/services/snippet_dsl_service.py @@ -32,7 +32,6 @@ from services.entities.dsl_entities import ( ) from services.plugin.dependencies_analysis import DependenciesAnalysisService from services.snippet_service import SNIPPET_FORBIDDEN_NODE_TYPES, SnippetService -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -491,29 +490,34 @@ class SnippetDslService: self._session.commit() if workflow_data: - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id=snippet.tenant_id, agent_ids=retirement_candidates, account_id=account.id, ) - enqueue_agent_resource_collection( - tenant_id=snippet.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) return snippet - def export_snippet_dsl(self, snippet: CustomizedSnippet, include_secret: bool = False) -> str: + def export_snippet_dsl( + self, snippet: CustomizedSnippet, include_secret: bool = False, workflow_id: str | None = None + ) -> str: """ Export snippet as DSL :param snippet: CustomizedSnippet instance :param include_secret: Whether include secret variable + :param workflow_id: Optional published workflow version to export; defaults to the draft workflow :return: YAML string """ snippet_service = self._snippet_service() - workflow = snippet_service.get_draft_workflow(snippet=snippet) + workflow = ( + snippet_service.get_published_workflow_by_id(snippet=snippet, workflow_id=workflow_id) + if workflow_id + else snippet_service.get_draft_workflow(snippet=snippet) + ) if not workflow: - raise ValueError("Missing draft workflow configuration, please check.") + workflow_description = ( + f"published workflow {workflow_id}" if workflow_id else "draft workflow configuration" + ) + raise ValueError(f"Missing {workflow_description}, please check.") icon_info = snippet.icon_info or {} export_data = { diff --git a/api/services/snippet_service.py b/api/services/snippet_service.py index f34f9789fa5..cf0809dd520 100644 --- a/api/services/snippet_service.py +++ b/api/services/snippet_service.py @@ -19,10 +19,11 @@ from models.agent import ( Agent, AgentScope, AgentStatus, + WorkflowAgentBindingType, WorkflowAgentNodeBinding, ) from models.enums import WorkflowRunTriggeredFrom -from models.model import App, AppMode, UploadFile +from models.model import UploadFile from models.snippet import CustomizedSnippet, SnippetType from models.tools import WorkflowToolProvider from models.workflow import ( @@ -39,13 +40,13 @@ from models.workflow import ( from repositories.factory import DifyAPIRepositoryFactory from services.agent.retirement_service import WorkflowAgentRetirementService from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError +from services.errors.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError from services.tag_service import TagService from services.workflow_node_execution_trace_service import ( WorkflowNodeExecutionTrace, assemble_workflow_node_execution_traces, ) from services.workflow_restore import apply_published_workflow_snapshot_to_draft -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -360,51 +361,51 @@ class SnippetService: snippet: CustomizedSnippet, account_id: str | None = None, ) -> bool: - """ - Delete a snippet. + """Stage Snippet deletion in the caller's transaction. + + Workflow rows and all of their binding owners are deleted in that + transaction. A single ``after_commit`` callback performs Agent + retirement, so rollback does not trigger cleanup. :param session: Database session :param snippet: Snippet to delete :return: True if deleted successfully """ SnippetService._delete_draft_variable_files(session=session, snippet=snippet) - owned_agents = session.scalars( - select(Agent).where( - Agent.tenant_id == snippet.tenant_id, - Agent.app_id == snippet.id, - Agent.scope == AgentScope.WORKFLOW_ONLY, - Agent.source.in_(WORKFLOW_ONLY_AGENT_SOURCES), - Agent.status == AgentStatus.ACTIVE, - ) - ).all() - now = datetime.now(UTC).replace(tzinfo=None) - backing_app_ids = {agent.backing_app_id for agent in owned_agents if agent.backing_app_id} - for agent in owned_agents: - agent.status = AgentStatus.ARCHIVED - agent.archived_by = account_id - agent.archived_at = now - agent.updated_by = account_id or agent.updated_by - agent.updated_at = now - - if backing_app_ids: - session.execute( - delete(App) - .where( - App.tenant_id == snippet.tenant_id, - App.id.in_(backing_app_ids), - App.mode == AppMode.AGENT, + candidate_agent_ids = { + agent_id + for agent_id in session.scalars( + select(WorkflowAgentNodeBinding.agent_id).where( + WorkflowAgentNodeBinding.tenant_id == snippet.tenant_id, + WorkflowAgentNodeBinding.app_id == snippet.id, + WorkflowAgentNodeBinding.binding_type == WorkflowAgentBindingType.INLINE_AGENT, + WorkflowAgentNodeBinding.agent_id.is_not(None), ) - .execution_options(synchronize_session=False) - ) + ).all() + if agent_id + } + candidate_agent_ids.update( + session.scalars( + select(Agent.id).where( + Agent.tenant_id == snippet.tenant_id, + Agent.app_id == snippet.id, + Agent.scope == AgentScope.WORKFLOW_ONLY, + Agent.source.in_(WORKFLOW_ONLY_AGENT_SOURCES), + Agent.status == AgentStatus.ACTIVE, + ) + ).all() + ) + if candidate_agent_ids: tenant_id = snippet.tenant_id - def cleanup_backing_apps(_session: Session) -> None: - from tasks.remove_app_and_related_data_task import remove_app_and_related_data_task + def collect_agent_resources(_session: Session) -> None: + WorkflowAgentRetirementService.retire_unowned( + tenant_id=tenant_id, + agent_ids=candidate_agent_ids, + account_id=account_id, + ) - for app_id in backing_app_ids: - remove_app_and_related_data_task.delay(tenant_id=tenant_id, app_id=app_id) - - event.listen(session, "after_commit", cleanup_backing_apps, once=True) + event.listen(session, "after_commit", collect_agent_resources, once=True) session.execute( delete(WorkflowAgentNodeBinding) @@ -620,16 +621,11 @@ class SnippetService: ) self._commit_if_owned(session) if self._session is None: - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id=snippet.tenant_id, agent_ids=retirement_candidates, account_id=account.id, ) - enqueue_agent_resource_collection( - tenant_id=snippet.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) return workflow def restore_published_workflow_to_draft( @@ -679,16 +675,11 @@ class SnippetService: ) self._commit_if_owned(session) if self._session is None: - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id=snippet.tenant_id, agent_ids=retirement_candidates, account_id=account.id, ) - enqueue_agent_resource_collection( - tenant_id=snippet.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) return draft_workflow def publish_workflow( @@ -697,7 +688,7 @@ class SnippetService: session: Session, snippet: CustomizedSnippet, account: Account, - ) -> tuple[Workflow, set[str]]: + ) -> Workflow: """ Publish the draft workflow as a new version. @@ -748,7 +739,7 @@ class SnippetService: kind=WorkflowKind.SNIPPET.value, ) session.add(workflow) - retirement_candidates = WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + WorkflowAgentPublishService.copy_agent_node_bindings_to_published( session=session, draft_workflow=draft_workflow, published_workflow=workflow, @@ -761,7 +752,7 @@ class SnippetService: snippet.updated_by = account.id session.add(snippet) - return workflow, retirement_candidates + return workflow def get_all_published_workflows( self, @@ -842,6 +833,53 @@ class SnippetService: session.add(workflow) return workflow + def delete_workflow( + self, + *, + session: Session, + snippet: CustomizedSnippet, + workflow_id: str, + ) -> bool: + """ + Delete a published snippet workflow version. + + :param session: Database session + :param snippet: CustomizedSnippet instance + :param workflow_id: Workflow ID + :return: True if successful + :raises: ValueError if workflow not found + :raises: WorkflowInUseError if workflow is the snippet's active version or published as a tool + :raises: DraftWorkflowDeletionError if workflow is a draft version + """ + stmt = select(Workflow).where( + Workflow.id == workflow_id, + Workflow.tenant_id == snippet.tenant_id, + Workflow.app_id == snippet.id, + self._snippet_kind_filter(), + ) + workflow = session.scalar(stmt) + if not workflow: + raise ValueError(f"Workflow with ID {workflow_id} not found") + + if workflow.version == Workflow.VERSION_DRAFT: + raise DraftWorkflowDeletionError("Cannot delete draft workflow versions") + + if snippet.workflow_id == workflow.id: + raise WorkflowInUseError(f"Cannot delete workflow that is currently in use by snippet '{snippet.id}'") + + tool_provider = session.scalar( + select(WorkflowToolProvider).where( + WorkflowToolProvider.tenant_id == snippet.tenant_id, + WorkflowToolProvider.app_id == snippet.id, + WorkflowToolProvider.version == workflow.version, + ) + ) + if tool_provider: + raise WorkflowInUseError("Cannot delete workflow that is published as a tool") + + session.delete(workflow) + return True + # --- Default Block Configs --- def get_default_block_configs(self) -> list[dict]: diff --git a/api/services/tag_application_service.py b/api/services/tag_application_service.py new file mode 100644 index 00000000000..04b4a931161 --- /dev/null +++ b/api/services/tag_application_service.py @@ -0,0 +1,103 @@ +"""Application boundary for Console tag management.""" + +from collections.abc import Sequence +from typing import Literal, NamedTuple, Protocol + +from machinery.context import RequestContext + +type TagKind = Literal["knowledge", "app", "snippet"] + + +class TagSummary(NamedTuple): + id: str + name: str + type: str + binding_count: int + + +class CreateTagInput(NamedTuple): + name: str + type: TagKind + + +class UpdateTagInput(NamedTuple): + name: str + + +class TagBindingInput(NamedTuple): + tag_ids: tuple[str, ...] + target_id: str + type: TagKind + + +class TagStore(Protocol): + def list_tags(self, workspace_id: str, tag_type: str, keyword: str | None) -> Sequence[TagSummary]: ... + + def get_tag_type(self, workspace_id: str, tag_id: str) -> str | None: ... + + def create_tag(self, workspace_id: str, actor_id: str, tag: CreateTagInput) -> TagSummary: ... + + def update_tag(self, workspace_id: str, tag_id: str, tag: UpdateTagInput) -> TagSummary: ... + + def delete_tag(self, workspace_id: str, tag_id: str) -> None: ... + + def create_bindings(self, workspace_id: str, actor_id: str, binding: TagBindingInput) -> None: ... + + def delete_bindings(self, workspace_id: str, binding: TagBindingInput) -> None: ... + + +class TagApplicationError(Exception): + """Base class for framework-neutral tag failures.""" + + +class TagNotFoundError(TagApplicationError): + def __init__(self) -> None: + super().__init__("Tag not found") + + +class TagNameConflictError(TagApplicationError): + def __init__(self) -> None: + super().__init__("Tag name already exists") + + +class TagBindingTargetNotFoundError(TagApplicationError): + def __init__(self, target_type: TagKind) -> None: + target_name = {"knowledge": "Dataset", "app": "App", "snippet": "Snippet"}[target_type] + super().__init__(f"{target_name} not found") + + +class InvalidTagBindingTypeError(TagApplicationError): + def __init__(self) -> None: + super().__init__("Invalid binding type") + + +class TagApplicationService: + def __init__(self, *, tags: TagStore) -> None: + self._tags = tags + + def list_tags(self, context: RequestContext, tag_type: str, keyword: str | None = None) -> tuple[TagSummary, ...]: + return tuple(self._tags.list_tags(self._workspace_id(context), tag_type, keyword)) + + def get_tag_type(self, context: RequestContext, tag_id: str) -> str | None: + return self._tags.get_tag_type(self._workspace_id(context), tag_id) + + def create_tag(self, context: RequestContext, tag: CreateTagInput) -> TagSummary: + return self._tags.create_tag(self._workspace_id(context), context.account_id, tag) + + def update_tag(self, context: RequestContext, tag_id: str, tag: UpdateTagInput) -> TagSummary: + return self._tags.update_tag(self._workspace_id(context), tag_id, tag) + + def delete_tag(self, context: RequestContext, tag_id: str) -> None: + self._tags.delete_tag(self._workspace_id(context), tag_id) + + def create_bindings(self, context: RequestContext, binding: TagBindingInput) -> None: + self._tags.create_bindings(self._workspace_id(context), context.account_id, binding) + + def delete_bindings(self, context: RequestContext, binding: TagBindingInput) -> None: + self._tags.delete_bindings(self._workspace_id(context), binding) + + @staticmethod + def _workspace_id(context: RequestContext) -> str: + if context.active_workspace_id is None: + raise RuntimeError("Console account admission did not resolve an active workspace") + return context.active_workspace_id diff --git a/api/services/trial_app_usage.py b/api/services/trial_app_usage.py new file mode 100644 index 00000000000..53cdb531f32 --- /dev/null +++ b/api/services/trial_app_usage.py @@ -0,0 +1,7 @@ +"""Port for recording recommended trial app usage.""" + +from typing import Protocol + + +class TrialAppUsageRecorder(Protocol): + def record(self, *, app_id: str, account_id: str) -> None: ... diff --git a/api/services/web_app_runtime_query_service.py b/api/services/web_app_runtime_query_service.py new file mode 100644 index 00000000000..559b0685d85 --- /dev/null +++ b/api/services/web_app_runtime_query_service.py @@ -0,0 +1,102 @@ +"""Application service for building the public Web app runtime bootstrap.""" + +import json +from collections.abc import Callable, Mapping +from typing import NamedTuple, Protocol, cast + +from services.app_definition_query_service import AppSiteConfiguration +from services.entities.feature_entities import FeatureModel +from services.file_service import FileService + + +class WebAppRuntimeRecord(NamedTuple): + app_id: str + tenant_id: str + mode: str + enable_site: bool + site: AppSiteConfiguration + plan: str + tenant_status: str + # Keep this lazy: workspaces without custom branding never parsed this legacy field. + tenant_custom_config_json: str | None + + +class WebAppBootstrap(NamedTuple): + app_id: str + mode: str + enable_site: bool + site: Mapping[str, str | bool | None] + plan: str + can_replace_logo: bool + custom_config: Mapping[str, str | bool | None] | None + + +class WebAppRuntimeQuery(Protocol): + def get_runtime_record(self, app_id: str) -> WebAppRuntimeRecord | None: ... + + +class WebAppRuntimeUnavailableError(ValueError): + """Raised when the admitted Web app can no longer be bootstrapped.""" + + +_ARCHIVED_TENANT_STATUS = "archive" + + +class WebAppRuntimeQueryService: + def __init__( + self, + *, + runtime: WebAppRuntimeQuery, + file_service: FileService, + workspace_features: Callable[[str], FeatureModel], + files_url: str, + ) -> None: + self._runtime = runtime + self._file_service = file_service + self._workspace_features = workspace_features + self._files_url = files_url + + def get_bootstrap(self, app_id: str) -> WebAppBootstrap: + record = self._runtime.get_runtime_record(app_id) + if record is None or record.tenant_status == _ARCHIVED_TENANT_STATUS: + raise WebAppRuntimeUnavailableError("Site not found") + + features = self._workspace_features(record.tenant_id) + site_icon_url = ( + self._file_service.get_icon_url(record.site.icon, record.tenant_id) + if record.site.icon_type == "image" and record.site.icon + else None + ) + + site = cast(dict[str, str | bool | None], record.site._asdict()) + site["icon_url"] = site_icon_url + if features.billing.enabled and not features.webapp_copyright_enabled: + site["copyright"] = None + site["input_placeholder"] = None + + custom_config = None + if features.can_replace_logo: + tenant_custom_config = ( + cast(Mapping[str, str | bool | None], json.loads(record.tenant_custom_config_json)) + if record.tenant_custom_config_json + else {} + ) + replace_webapp_logo = ( + f"{self._files_url}/files/workspaces/{record.tenant_id}/webapp-logo" + if tenant_custom_config.get("replace_webapp_logo") + else None + ) + custom_config = { + "remove_webapp_brand": tenant_custom_config.get("remove_webapp_brand", False), + "replace_webapp_logo": replace_webapp_logo, + } + + return WebAppBootstrap( + app_id=record.app_id, + mode=record.mode, + enable_site=record.enable_site, + site=site, + plan=record.plan, + can_replace_logo=features.can_replace_logo, + custom_config=custom_config, + ) diff --git a/api/services/webapp_access_query_service.py b/api/services/webapp_access_query_service.py index 10524041f9c..1228d0affbd 100644 --- a/api/services/webapp_access_query_service.py +++ b/api/services/webapp_access_query_service.py @@ -5,6 +5,8 @@ from typing import Protocol from enums import WebAppAccessMode +_PERMISSION_CHECK_MODES = frozenset({WebAppAccessMode.PRIVATE, WebAppAccessMode.PRIVATE_ALL}) + class WebAppAccessQuery(Protocol): def find_app_id_by_code(self, app_code: str) -> str | None: ... @@ -29,10 +31,12 @@ class WebAppAccessQueryService: access: WebAppAccessQuery, webapp_auth_enabled: bool, access_mode_for_app: Callable[[str], WebAppAccessMode], + is_user_allowed_for_app: Callable[[str, str], bool], ) -> None: self._access = access self._webapp_auth_enabled = webapp_auth_enabled self._access_mode_for_app = access_mode_for_app + self._is_user_allowed_for_app = is_user_allowed_for_app def get_access_mode(self, *, app_id: str | None, app_code: str | None) -> WebAppAccessMode: if not self._webapp_auth_enabled: @@ -47,3 +51,12 @@ class WebAppAccessQueryService: raise WebAppAccessReferenceRequiredError("appId or appCode must be provided") return self._access_mode_for_app(app_id) + + def requires_permission_check(self, app_id: str) -> bool: + return self._access_mode_for_app(app_id) in _PERMISSION_CHECK_MODES + + def is_user_allowed(self, *, user_id: str, app_id: str) -> bool: + if not self._webapp_auth_enabled: + return True + + return self._is_user_allowed_for_app(user_id, app_id) diff --git a/api/services/workflow_collaboration_service.py b/api/services/workflow_collaboration_service.py index 11fae8ebbf6..8e1f0c8a01a 100644 --- a/api/services/workflow_collaboration_service.py +++ b/api/services/workflow_collaboration_service.py @@ -12,9 +12,12 @@ from socketio.exceptions import TimeoutError as SocketIOTimeoutError # type: ig from sqlalchemy import select from sqlalchemy.orm import Session +from configs import dify_config +from core.rbac import RBACPermission, RBACResourceScope from models.account import Account from models.model import App from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository, WorkflowSessionInfo +from services.enterprise.rbac_service import RBACService logger = logging.getLogger(__name__) @@ -112,7 +115,7 @@ class WorkflowCollaborationService: if not user_id or not tenant_id: return None - if not self._can_access_workflow(workflow_id, str(tenant_id), session=session): + if not self._can_access_workflow(workflow_id, str(tenant_id), str(user_id), session=session): logger.warning( "Workflow collaboration join rejected: workflow_id=%s tenant_id=%s user_id=%s sid=%s", workflow_id, @@ -148,10 +151,27 @@ class WorkflowCollaborationService: return str(user_id), is_leader - def _can_access_workflow(self, workflow_id: str, tenant_id: str, *, session: Session) -> bool: - """Check room access without relying on Flask's app-context-bound scoped session.""" - app_id = session.scalar(select(App.id).where(App.id == workflow_id, App.tenant_id == tenant_id).limit(1)) - return app_id is not None + def _can_access_workflow(self, workflow_id: str, tenant_id: str, user_id: str, *, session: Session) -> bool: + """Check tenant and app permission without relying on Flask's scoped session.""" + with session.begin(): + app = session.execute( + select(App.id, App.maintainer).where( + App.id == workflow_id, App.tenant_id == tenant_id, App.status == "normal" + ) + ).one_or_none() + if app is None: + return False + + app_id, maintainer = app + if not dify_config.RBAC_ENABLED or maintainer == user_id: + return True + return RBACService.CheckAccess.check( + tenant_id, + user_id, + scene=RBACPermission.APP_EDIT, + resource_type=RBACResourceScope.APP, + resource_id=app_id, + ) def disconnect_session(self, sid: str) -> None: mapping = self._repository.get_sid_mapping(sid) diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index 2b673ec5cc1..157290aaca2 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -80,6 +80,7 @@ from graphon.variables.input_entities import VariableEntityType from graphon.variables.variables import Variable from libs.datetime_utils import naive_utc_now from models import Account +from models.agent import WorkflowAgentBindingType, WorkflowAgentNodeBinding from models.human_input import HumanInputFormRecipient, RecipientType from models.model import App, AppMode from models.tools import WorkflowToolProvider @@ -93,7 +94,6 @@ from services.errors.app import ( WorkflowHashNotEqualError, WorkflowNotFoundError, ) -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection @dataclass(frozen=True) @@ -280,14 +280,21 @@ class WorkflowService: .with_for_update() ) - def get_published_workflow_by_id(self, app_model: App, workflow_id: str, *, session: Session) -> Workflow | None: - """ - fetch published workflow by workflow_id + def get_published_workflow_by_id( + self, + app_model: App, + workflow_id: str, + *, + session: Session, + for_update: bool = False, + ) -> Workflow | None: + """Fetch a published workflow by ID in the caller's transaction. - Reuses the caller's active session so workflow reads stay in the same - transaction as the surrounding request or task. + With ``for_update=True``, the source version stays locked until that + transaction ends. Restore uses the lock while copying Agent bindings so + a concurrent delete cannot release the same owner. """ - workflow = session.scalar( + stmt = ( select(Workflow) .where( Workflow.tenant_id == app_model.tenant_id, @@ -296,6 +303,9 @@ class WorkflowService: ) .limit(1) ) + if for_update: + stmt = stmt.with_for_update() + workflow = session.scalar(stmt) if not workflow: return None if workflow.version == Workflow.VERSION_DRAFT: @@ -328,15 +338,17 @@ class WorkflowService: return workflow - def get_accessible_app_ids(self, app_ids: Sequence[str], tenant_id: str, *, session: Session) -> set[str]: - """ - Return app IDs that belong to the given tenant. - """ + def get_tenant_app_maintainers( + self, app_ids: Sequence[str], tenant_id: str, *, session: Session + ) -> dict[str, str | None]: + """Return requested normal apps and their maintainers within a tenant.""" if not app_ids: - return set() + return {} - stmt = select(App.id).where(App.id.in_(app_ids), App.tenant_id == tenant_id) - return {str(app_id) for app_id in session.scalars(stmt).all()} + stmt = select(App.id, App.maintainer).where( + App.id.in_(app_ids), App.tenant_id == tenant_id, App.status == "normal" + ) + return {str(app_id): maintainer for app_id, maintainer in session.execute(stmt)} def get_all_published_workflow( self, @@ -491,16 +503,11 @@ class WorkflowService: # commit db session changes if commit: session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id=app_model.tenant_id, agent_ids=retirement_candidates, account_id=account.id, ) - enqueue_agent_resource_collection( - tenant_id=app_model.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) # trigger app workflow events if commit: @@ -624,7 +631,10 @@ class WorkflowService: published workflow so the normal draft sync flow stays stateless. """ source_workflow = self.get_published_workflow_by_id( - app_model=app_model, workflow_id=workflow_id, session=session + app_model=app_model, + workflow_id=workflow_id, + session=session, + for_update=True, ) if not source_workflow: raise WorkflowNotFoundError("Workflow not found.") @@ -656,16 +666,11 @@ class WorkflowService: ) session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id=app_model.tenant_id, agent_ids=retirement_candidates, account_id=account.id, ) - enqueue_agent_resource_collection( - tenant_id=app_model.tenant_id, - binding_ids=binding_ids, - home_snapshot_ids=home_snapshot_ids, - ) app_draft_workflow_was_synced.send(app_model, synced_draft_workflow=draft_workflow) return draft_workflow @@ -678,7 +683,7 @@ class WorkflowService: account: Account, marked_name: str = "", marked_comment: str = "", - ) -> tuple[Workflow, set[str]]: + ) -> Workflow: draft_workflow_stmt = select(Workflow).where( Workflow.tenant_id == app_model.tenant_id, Workflow.app_id == app_model.id, @@ -752,7 +757,7 @@ class WorkflowService: # commit db session changes session.add(workflow) - retirement_candidates = WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + WorkflowAgentPublishService.copy_agent_node_bindings_to_published( session=session, draft_workflow=draft_workflow, published_workflow=workflow, @@ -766,7 +771,7 @@ class WorkflowService: ) # return new workflow - return workflow, retirement_candidates + return workflow def _validate_workflow_credentials(self, workflow: Workflow, *, session: Session) -> None: """ @@ -1901,21 +1906,29 @@ class WorkflowService: return workflow - def delete_workflow(self, *, session: Session, workflow_ref: WorkflowRef) -> bool: - """ - Delete a workflow + def delete_workflow(self, *, session: Session, workflow_ref: WorkflowRef) -> list[str]: + """Stage a published Workflow and its binding owners for deletion. + + The exact owner key is tenant, App, Workflow, and Workflow version. The + Workflow row lock serializes source-version reads and restoration with + deletion. The caller must commit successfully before retiring the + returned, sorted and deduplicated inline Agent candidates. :param session: SQLAlchemy database session :param workflow_ref: Owner-bound workflow reference - :return: True if successful + :return: Inline Agent IDs whose owner binding is staged for deletion :raises: ValueError if workflow not found :raises: WorkflowInUseError if workflow is in use :raises: DraftWorkflowDeletionError if workflow is a draft version """ - stmt = select(Workflow).where( - Workflow.id == workflow_ref.workflow_id, - Workflow.tenant_id == workflow_ref.tenant_id, - Workflow.app_id == workflow_ref.owner_id, + stmt = ( + select(Workflow) + .where( + Workflow.id == workflow_ref.workflow_id, + Workflow.tenant_id == workflow_ref.tenant_id, + Workflow.app_id == workflow_ref.owner_id, + ) + .with_for_update() ) workflow = session.scalar(stmt) @@ -1947,8 +1960,25 @@ class WorkflowService: # Cannot delete a workflow that's published as a tool raise WorkflowInUseError("Cannot delete workflow that is published as a tool") + bindings = session.scalars( + select(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.tenant_id == workflow.tenant_id, + WorkflowAgentNodeBinding.app_id == workflow.app_id, + WorkflowAgentNodeBinding.workflow_id == workflow.id, + WorkflowAgentNodeBinding.workflow_version == workflow.version, + ) + ).all() + retirement_candidates = sorted( + { + binding.agent_id + for binding in bindings + if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id + } + ) + for binding in bindings: + session.delete(binding) session.delete(workflow) - return True + return retirement_candidates def _setup_variable_pool( diff --git a/api/tasks/collect_agent_resources_task.py b/api/tasks/collect_agent_resources_task.py index bfec87b7bf3..54addc2f45a 100644 --- a/api/tasks/collect_agent_resources_task.py +++ b/api/tasks/collect_agent_resources_task.py @@ -1,4 +1,10 @@ -"""Asynchronously collect retired Agent working resources.""" +"""Collect retired Agent data under a two-phase task contract. + +Phase one attempts every explicitly identified RETIRED working resource. Phase +two purges the requested archived Agent aggregates only after the whole first +phase succeeds. Any collection failure skips aggregate purge and is included in +the error raised after all explicit resources have been attempted. +""" from __future__ import annotations @@ -7,6 +13,7 @@ from collections.abc import Iterable from celery import shared_task +from services.agent.deletion_service import AgentDeletionService from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.workspace_service import AgentWorkspaceService @@ -20,8 +27,14 @@ def collect_agent_resources( binding_ids: list[str], workspace_ids: list[str], home_snapshot_ids: list[str], + purge_agent_ids: list[str] | None = None, ) -> None: - """Collect only the explicitly identified RETIRED resources.""" + """Collect the explicit RETIRED batch, then purge Agents only on full success. + + Collection is best-effort across the complete explicit batch so one failed + resource does not hide later failures. If any resource fails, aggregate + purge is skipped and one summary error is raised after all attempts. + """ collectors = ( (workspace_ids, "workspace_id", AgentWorkspaceService.collect_retired_workspace), @@ -32,20 +45,30 @@ def collect_agent_resources( AgentHomeSnapshotService.collect_retired_home_snapshot, ), ) + failures: list[str] = [] + first_error: Exception | None = None for resource_ids, argument_name, collector in collectors: for resource_id in resource_ids: try: collector(tenant_id=tenant_id, **{argument_name: resource_id}) - except Exception: + except Exception as exc: + resource_type = argument_name.removesuffix("_id") + failures.append(f"{resource_type}:{resource_id}") + if first_error is None: + first_error = exc logger.exception( "Failed to collect retired Agent resource", extra={ "tenant_id": tenant_id, - "resource_type": argument_name.removesuffix("_id"), + "resource_type": resource_type, "resource_id": resource_id, }, ) - raise + if failures: + raise RuntimeError( + f"Failed to collect {len(failures)} retired Agent resource(s): {', '.join(failures)}" + ) from first_error + AgentDeletionService.purge_archived_agents(tenant_id=tenant_id, agent_ids=purge_agent_ids or ()) def enqueue_agent_resource_collection( @@ -54,13 +77,15 @@ def enqueue_agent_resource_collection( binding_ids: Iterable[str] = (), workspace_ids: Iterable[str] = (), home_snapshot_ids: Iterable[str] = (), + purge_agent_ids: Iterable[str] = (), ) -> None: - """Best-effort enqueue of physical collection after retire has committed.""" + """Enqueue physical collection after retirement has committed.""" payload = { "binding_ids": sorted({resource_id for resource_id in binding_ids if resource_id}), "workspace_ids": sorted({resource_id for resource_id in workspace_ids if resource_id}), "home_snapshot_ids": sorted({resource_id for resource_id in home_snapshot_ids if resource_id}), + "purge_agent_ids": sorted({agent_id for agent_id in purge_agent_ids if agent_id}), } if not any(payload.values()): return @@ -71,6 +96,7 @@ def enqueue_agent_resource_collection( "Failed to enqueue retired Agent resource collection", extra={"tenant_id": tenant_id, **payload}, ) + raise __all__ = ["collect_agent_resources", "enqueue_agent_resource_collection"] diff --git a/api/tasks/remove_app_and_related_data_task.py b/api/tasks/remove_app_and_related_data_task.py index 4562b7d1d90..f8af48e0990 100644 --- a/api/tasks/remove_app_and_related_data_task.py +++ b/api/tasks/remove_app_and_related_data_task.py @@ -40,6 +40,7 @@ from models import ( TraceAppConfig, WorkflowSchedulePlan, ) +from models.agent import WorkflowAgentNodeBinding from models.tools import WorkflowToolProvider from models.trigger import WorkflowPluginTrigger, WorkflowTriggerLog, WorkflowWebhookTrigger from models.web import PinnedConversation, SavedMessage @@ -70,6 +71,7 @@ def remove_app_and_related_data_task(self, tenant_id: str, app_id: str): _delete_recommended_apps(tenant_id, app_id) _delete_app_annotation_data(tenant_id, app_id) _delete_app_dataset_joins(tenant_id, app_id) + _delete_workflow_agent_node_bindings(tenant_id, app_id) _delete_app_workflows(tenant_id, app_id) _delete_app_workflow_runs(tenant_id, app_id) _delete_app_workflow_node_executions(tenant_id, app_id) @@ -262,6 +264,17 @@ def _delete_app_workflows(tenant_id: str, app_id: str): ) +def _delete_workflow_agent_node_bindings(tenant_id: str, app_id: str) -> None: + with session_factory.create_session() as session: + session.execute( + delete(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.tenant_id == tenant_id, + WorkflowAgentNodeBinding.app_id == app_id, + ) + ) + session.commit() + + def _delete_app_workflow_runs(tenant_id: str, app_id: str): """Delete all workflow runs for an app using the service repository.""" session_maker = sessionmaker(bind=db.engine) diff --git a/api/tests/test_containers_integration_tests/.ruff.toml b/api/tests/test_containers_integration_tests/.ruff.toml index 49244091ffb..a58180c68b1 100644 --- a/api/tests/test_containers_integration_tests/.ruff.toml +++ b/api/tests/test_containers_integration_tests/.ruff.toml @@ -23,7 +23,6 @@ extend-select = ["ANN401", "ARG"] "services/dataset_collection_binding.py" = ["ARG002"] "services/document_service_status.py" = ["ARG002"] "services/rag_pipeline/test_rag_pipeline_service_db.py" = ["ARG002"] -"services/recommend_app/test_database_retrieval.py" = ["ARG002"] "services/test_account_service.py" = ["ARG002"] "services/test_advanced_prompt_template_service.py" = ["ARG002"] "services/test_app_dsl_service.py" = ["ANN401", "ARG001", "ARG002", "ARG005", "TID251"] diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_human_input_form.py b/api/tests/test_containers_integration_tests/controllers/web/test_human_input_form.py index 5c3ad5e199b..d26965a0e99 100644 --- a/api/tests/test_containers_integration_tests/controllers/web/test_human_input_form.py +++ b/api/tests/test_containers_integration_tests/controllers/web/test_human_input_form.py @@ -231,7 +231,7 @@ def test_get_human_input_form_resolves_runtime_select_options( return features monkeypatch.setattr( - "controllers.web.site.FeatureService.get_features", + "controllers.web.human_input_form.FeatureService.get_features", mock_get_features, ) diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_site.py b/api/tests/test_containers_integration_tests/controllers/web/test_site.py index cdba83851fc..e5f6f83d768 100644 --- a/api/tests/test_containers_integration_tests/controllers/web/test_site.py +++ b/api/tests/test_containers_integration_tests/controllers/web/test_site.py @@ -9,10 +9,7 @@ from flask import Flask from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden -from configs import dify_config from controllers.web.site import AppSiteApi, WebAppSiteResponse, WebModelConfigResponse -from enums import DeploymentEdition -from extensions.storage.storage_type import StorageType from models import Tenant, TenantStatus from models.account import TenantCustomConfigDict from models.model import App, AppMode, AppModelConfig, CustomizeTokenStrategy, EndUser, Site @@ -82,7 +79,7 @@ def _site_model(*, app_id: str) -> Site: class TestAppSiteApi: - @patch("controllers.web.site.FeatureService.get_features") + @patch("services.feature_service.FeatureService.get_features") def test_happy_path(self, mock_features: MagicMock, app: Flask, db_session_with_containers: Session) -> None: app.config["RESTX_MASK_HEADER"] = "X-Fields" tenant = _create_tenant(db_session_with_containers) @@ -100,39 +97,6 @@ class TestAppSiteApi: assert result["enable_site"] is True assert result["mode"] == AppMode.CHAT - @patch("controllers.web.site.FileService.get_file_presigned_url") - @patch("controllers.web.site.FeatureService.get_features") - def test_image_icon_uses_s3_presigned_url( - self, - mock_features: MagicMock, - mock_get_file_presigned_url: MagicMock, - app: Flask, - db_session_with_containers: Session, - ) -> None: - app.config["RESTX_MASK_HEADER"] = "X-Fields" - tenant = _create_tenant(db_session_with_containers) - app_model = _create_app(db_session_with_containers, tenant.id) - site = _create_site(db_session_with_containers, app_model.id) - site.icon_type = "image" - site.icon = "11111111-1111-4111-8111-111111111111" - db_session_with_containers.commit() - end_user = _end_user(tenant.id, app_model.id) - mock_features.return_value = FeatureModel(can_replace_logo=False) - mock_get_file_presigned_url.return_value = "https://s3.example.com/icon.png?signature=test" - - with ( - patch.object(dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.CLOUD), - patch.object(dify_config, "STORAGE_TYPE", StorageType.S3), - app.test_request_context("/site"), - ): - result = AppSiteApi().get(app_model, end_user) - - assert result["site"]["icon_url"] == "https://s3.example.com/icon.png?signature=test" - mock_get_file_presigned_url.assert_called_once_with( - file_id="11111111-1111-4111-8111-111111111111", - tenant_id=tenant.id, - ) - def test_missing_site_raises_forbidden(self, app: Flask, db_session_with_containers: Session) -> None: app.config["RESTX_MASK_HEADER"] = "X-Fields" tenant = _create_tenant(db_session_with_containers) diff --git a/api/tests/test_containers_integration_tests/pyrefly.toml b/api/tests/test_containers_integration_tests/pyrefly.toml index cf707c4f947..cd654925428 100644 --- a/api/tests/test_containers_integration_tests/pyrefly.toml +++ b/api/tests/test_containers_integration_tests/pyrefly.toml @@ -42,7 +42,6 @@ project-excludes = [ "services/dataset_collection_binding.py", "services/dataset_service_update_delete.py", "services/document_service_status.py", - "services/recommend_app/test_database_retrieval.py", "services/test_account_service.py", "services/test_advanced_prompt_template_service.py", "services/test_agent_service.py", diff --git a/api/tests/test_containers_integration_tests/repositories/test_recommended_app_catalog_repository.py b/api/tests/test_containers_integration_tests/repositories/test_recommended_app_catalog_repository.py new file mode 100644 index 00000000000..242d8c8324c --- /dev/null +++ b/api/tests/test_containers_integration_tests/repositories/test_recommended_app_catalog_repository.py @@ -0,0 +1,119 @@ +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from sqlalchemy.orm import Session, object_session, sessionmaker + +from extensions.ext_redis import RedisClientWrapper +from models.model import App, RecommendedApp, Site +from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository +from services.recommended_app_query_service import RecommendedAppDetailRecord + + +def _add_catalog_app( + session: Session, + *, + categories: list[str] | None = None, + language: str = "en-US", + is_public: bool = True, + with_site: bool = True, +) -> App: + app = App( + tenant_id=str(uuid4()), + name=f"app-{uuid4()}", + mode="chat", + enable_site=True, + enable_api=True, + is_public=is_public, + ) + app.id = str(uuid4()) + session.add(app) + session.add( + RecommendedApp( + app_id=app.id, + description={"en-US": "test"}, + copyright="copy", + privacy_policy="privacy", + category="writing", + categories=["writing"] if categories is None else categories, + language=language, + is_listed=True, + position=1, + ) + ) + if with_site: + session.add( + Site( + app_id=app.id, + title=f"site-{uuid4()}", + default_language="en-US", + customize_token_strategy="not_allow", + description="description", + copyright="copyright", + privacy_policy="privacy", + custom_disclaimer="disclaimer", + ) + ) + session.commit() + return app + + +def _repository(session: Session) -> DatabaseRecommendedAppCatalogRepository: + redis = MagicMock(spec=RedisClientWrapper) + redis.get.return_value = None + return DatabaseRecommendedAppCatalogRepository( + sessionmaker(bind=session.get_bind(), expire_on_commit=False), + redis=redis, + ) + + +def test_list_maps_postgres_models_with_owned_session( + db_session_with_containers: Session, +) -> None: + app = _add_catalog_app( + db_session_with_containers, + categories=["writing", "assistant"], + ) + private_app = _add_catalog_app(db_session_with_containers, is_public=False) + no_site_app = _add_catalog_app(db_session_with_containers, with_site=False) + + page = _repository(db_session_with_containers).list_recommended("fr-FR") + + record = next(item for item in page.recommended_apps if item.app_id == app.id) + assert record.app is not None + assert record.app.id == app.id + assert record.app.mode == "chat" + assert record.description == "description" + assert record.categories == ("writing", "assistant") + assert {"writing", "assistant"} <= set(page.categories) + assert private_app.id not in {item.app_id for item in page.recommended_apps} + assert no_site_app.id not in {item.app_id for item in page.recommended_apps} + + +def test_membership_does_not_export_dsl_with_owned_session( + db_session_with_containers: Session, +) -> None: + app = _add_catalog_app(db_session_with_containers, with_site=False) + repository = _repository(db_session_with_containers) + + def export_dsl(*, app_model: App, session: Session) -> str: + assert object_session(app_model) is session + assert session is not db_session_with_containers + return "exported_yaml" + + with patch( + "repositories.recommended_app_catalog_repository.AppDslService.export_dsl", + side_effect=export_dsl, + ) as mock_export_dsl: + detail = repository.get_detail(app.id) + is_in_catalog = repository.contains(app.id) + + assert detail == RecommendedAppDetailRecord( + id=app.id, + name=app.name, + icon=app.icon, + icon_background=app.icon_background, + mode="chat", + export_data="exported_yaml", + ) + assert is_in_catalog is True + mock_export_dsl.assert_called_once() diff --git a/api/tests/test_containers_integration_tests/services/recommend_app/__init__.py b/api/tests/test_containers_integration_tests/services/recommend_app/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/api/tests/test_containers_integration_tests/services/recommend_app/test_database_retrieval.py b/api/tests/test_containers_integration_tests/services/recommend_app/test_database_retrieval.py deleted file mode 100644 index 145dc947b52..00000000000 --- a/api/tests/test_containers_integration_tests/services/recommend_app/test_database_retrieval.py +++ /dev/null @@ -1,301 +0,0 @@ -from __future__ import annotations - -from unittest.mock import patch -from uuid import uuid4 - -from flask import Flask -from sqlalchemy.orm import Session - -from models.model import App, RecommendedApp, Site -from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval - - -def _create_app(db_session: Session, *, tenant_id: str, is_public: bool = True) -> App: - app = App( - tenant_id=tenant_id, - name=f"app-{uuid4()}", - mode="chat", - enable_site=True, - enable_api=True, - is_public=is_public, - ) - app.id = str(uuid4()) - db_session.add(app) - db_session.commit() - return app - - -def _create_site(db_session: Session, *, app_id: str) -> Site: - site = Site( - app_id=app_id, - title=f"site-{uuid4()}", - default_language="en-US", - customize_token_strategy="not_allow", - description="desc", - copyright="copy", - privacy_policy="pp", - custom_disclaimer="cd", - ) - site.id = str(uuid4()) - db_session.add(site) - db_session.commit() - return site - - -def _create_recommended_app( - db_session, - *, - app_id: str, - category: str = "chat", - categories: list[str] | None = None, - language: str = "en-US", - is_listed: bool = True, - is_learn_dify: bool = False, - position: int = 1, -) -> RecommendedApp: - rec = RecommendedApp( - app_id=app_id, - description={"en-US": "test"}, - copyright="copy", - privacy_policy="pp", - category=category, - categories=[category] if categories is None else categories, - language=language, - is_listed=is_listed, - is_learn_dify=is_learn_dify, - position=position, - ) - rec.id = str(uuid4()) - db_session.add(rec) - db_session.commit() - return rec - - -class TestFetchRecommendedAppsFromDb: - def test_returns_apps_and_sorted_categories( - self, flask_app_with_containers: Flask, db_session_with_containers: Session - ): - tenant_id = str(uuid4()) - app1 = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=app1.id) - _create_recommended_app(db_session_with_containers, app_id=app1.id, category="writing") - - app2 = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=app2.id) - _create_recommended_app(db_session_with_containers, app_id=app2.id, category="assistant") - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( - "en-US", session=db_session_with_containers - ) - - app_ids = {r["app_id"] for r in result["recommended_apps"]} - assert app1.id in app_ids - assert app2.id in app_ids - assert "assistant" in result["categories"] - assert "writing" in result["categories"] - - def test_returns_multiple_categories_for_one_app( - self, flask_app_with_containers: Flask, db_session_with_containers: Session - ): - tenant_id = str(uuid4()) - created_app = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=created_app.id) - _create_recommended_app( - db_session_with_containers, - app_id=created_app.id, - category="writing", - categories=["writing", "assistant"], - ) - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( - "en-US", session=db_session_with_containers - ) - - recommended_app = next(item for item in result["recommended_apps"] if item["app_id"] == created_app.id) - assert recommended_app["categories"] == ["writing", "assistant"] - assert "writing" in result["categories"] - assert "assistant" in result["categories"] - - def test_ignores_legacy_category_when_categories_are_empty( - self, - flask_app_with_containers: Flask, - db_session_with_containers: Session, - ): - legacy_category = f"legacy-empty-{uuid4()}" - tenant_id = str(uuid4()) - created_app = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=created_app.id) - _create_recommended_app( - db_session_with_containers, - app_id=created_app.id, - category=legacy_category, - categories=[], - ) - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( - "en-US", session=db_session_with_containers - ) - - recommended_app = next(item for item in result["recommended_apps"] if item["app_id"] == created_app.id) - assert "category" not in recommended_app - assert recommended_app["categories"] == [] - assert legacy_category not in result["categories"] - - def test_falls_back_to_default_language_when_empty( - self, flask_app_with_containers: Flask, db_session_with_containers: Session - ): - tenant_id = str(uuid4()) - app1 = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=app1.id) - _create_recommended_app(db_session_with_containers, app_id=app1.id, language="en-US") - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( - "fr-FR", session=db_session_with_containers - ) - - app_ids = {r["app_id"] for r in result["recommended_apps"]} - assert app1.id in app_ids - - def test_skips_non_public_apps(self, flask_app_with_containers: Flask, db_session_with_containers: Session): - tenant_id = str(uuid4()) - app1 = _create_app(db_session_with_containers, tenant_id=tenant_id, is_public=False) - _create_site(db_session_with_containers, app_id=app1.id) - _create_recommended_app(db_session_with_containers, app_id=app1.id) - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( - "en-US", session=db_session_with_containers - ) - - app_ids = {r["app_id"] for r in result["recommended_apps"]} - assert app1.id not in app_ids - - def test_skips_apps_without_site(self, flask_app_with_containers: Flask, db_session_with_containers: Session): - tenant_id = str(uuid4()) - app1 = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_recommended_app(db_session_with_containers, app_id=app1.id) - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( - "en-US", session=db_session_with_containers - ) - - app_ids = {r["app_id"] for r in result["recommended_apps"]} - assert app1.id not in app_ids - - def test_fetch_learn_dify_apps_uses_flag_not_categories( - self, - flask_app_with_containers, - db_session_with_containers: Session, - ): - tenant_id = str(uuid4()) - learn_dify_app = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=learn_dify_app.id) - _create_recommended_app( - db_session_with_containers, - app_id=learn_dify_app.id, - category="workflow", - categories=["Workflow"], - is_learn_dify=True, - ) - - category_only_app = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=category_only_app.id) - _create_recommended_app( - db_session_with_containers, - app_id=category_only_app.id, - category="Learn Dify", - categories=["Learn Dify"], - is_learn_dify=False, - ) - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db( - "en-US", session=db_session_with_containers - ) - - app_ids = {r["app_id"] for r in result["recommended_apps"]} - assert learn_dify_app.id in app_ids - assert category_only_app.id not in app_ids - recommended_app = next(item for item in result["recommended_apps"] if item["app_id"] == learn_dify_app.id) - assert recommended_app["categories"] == ["Workflow"] - - def test_fetch_learn_dify_apps_falls_back_to_default_language( - self, - flask_app_with_containers, - db_session_with_containers: Session, - ): - tenant_id = str(uuid4()) - learn_dify_app = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=learn_dify_app.id) - _create_recommended_app( - db_session_with_containers, - app_id=learn_dify_app.id, - categories=["Workflow"], - is_learn_dify=True, - language="en-US", - ) - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db( - "fr-FR", session=db_session_with_containers - ) - - app_ids = {r["app_id"] for r in result["recommended_apps"]} - assert learn_dify_app.id in app_ids - - -class TestFetchRecommendedAppDetailFromDb: - def test_returns_none_when_not_listed(self, flask_app_with_containers: Flask, db_session_with_containers: Session): - result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db( - str(uuid4()), session=db_session_with_containers - ) - - assert result is None - - def test_returns_none_when_app_not_public( - self, flask_app_with_containers: Flask, db_session_with_containers: Session - ): - tenant_id = str(uuid4()) - app1 = _create_app(db_session_with_containers, tenant_id=tenant_id, is_public=False) - _create_recommended_app(db_session_with_containers, app_id=app1.id) - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db( - app1.id, session=db_session_with_containers - ) - - assert result is None - - @patch("services.recommend_app.database.database_retrieval.AppDslService") - def test_returns_detail_on_success( - self, mock_dsl, flask_app_with_containers: Flask, db_session_with_containers: Session - ): - tenant_id = str(uuid4()) - app1 = _create_app(db_session_with_containers, tenant_id=tenant_id) - _create_site(db_session_with_containers, app_id=app1.id) - _create_recommended_app(db_session_with_containers, app_id=app1.id) - mock_dsl.export_dsl.return_value = "exported_yaml" - - db_session_with_containers.expire_all() - - result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db( - app1.id, session=db_session_with_containers - ) - - assert result is not None - assert result["id"] == app1.id - assert result["export_data"] == "exported_yaml" diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_service.py b/api/tests/test_containers_integration_tests/services/test_workflow_service.py index 05697997511..c9d8ee51c72 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workflow_service.py @@ -871,13 +871,12 @@ class TestWorkflowService: from unittest.mock import patch with patch("flask_login.utils._get_user", return_value=account, autospec=True): - result, retirement_candidates = workflow_service.publish_workflow( + result = workflow_service.publish_workflow( session=db_session_with_containers, app_model=app, account=account ) # Assert assert result is not None - assert retirement_candidates == set() assert result.version != Workflow.VERSION_DRAFT # Version should be a timestamp format like '2025-08-22 00:10:24.722051' assert isinstance(result.version, str) @@ -1435,7 +1434,7 @@ class TestWorkflowService: ) # Assert - assert result is True + assert result == [] # Verify workflow is actually deleted deleted_workflow = db_session_with_containers.query(Workflow).filter_by(id=workflow.id).first() diff --git a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_deletion.py b/api/tests/test_containers_integration_tests/services/workflow/test_workflow_deletion.py index d2fc873cf51..3c70dea134f 100644 --- a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_deletion.py +++ b/api/tests/test_containers_integration_tests/services/workflow/test_workflow_deletion.py @@ -116,7 +116,7 @@ class TestWorkflowDeletion: workflow_ref=WorkflowRef(tenant_id=tenant.id, owner_id=app.id, workflow_id=workflow_id), ) - assert result is True + assert result == [] db_session_with_containers.expire_all() assert db_session_with_containers.get(Workflow, workflow_id) is None diff --git a/api/tests/unit_tests/.ruff.toml b/api/tests/unit_tests/.ruff.toml index aa04b37fd7c..32c17b823bd 100644 --- a/api/tests/unit_tests/.ruff.toml +++ b/api/tests/unit_tests/.ruff.toml @@ -347,7 +347,6 @@ extend-select = ["ANN401", "ARG"] "services/rag_pipeline/test_rag_pipeline_service.py" = ["ARG001", "ARG005"] "services/rag_pipeline/test_rag_pipeline_task_proxy.py" = ["ARG001", "ARG005"] "services/rag_pipeline/test_rag_pipeline_transform_service.py" = ["ARG001"] -"services/recommend_app/test_remote_retrieval.py" = ["ARG002"] "services/retention/workflow_run/test_archive_download_preparation.py" = ["ARG002"] "services/retention/workflow_run/test_archive_log_service.py" = ["ARG001", "ARG002"] "services/retention/workflow_run/test_bundle_archive_maintenance.py" = ["TID251"] @@ -379,7 +378,6 @@ extend-select = ["ANN401", "ARG"] "services/test_oauth_server_service.py" = ["ARG002"] "services/test_operation_service.py" = ["TID251"] "services/test_rag_pipeline_task_proxy.py" = ["ARG002"] -"services/test_recommended_app_service.py" = ["ARG001"] "services/test_schedule_service.py" = ["ANN401", "TID251"] "services/test_snippet_service.py" = ["ARG001", "ARG002"] "services/test_summary_index_service.py" = ["ARG001"] diff --git a/api/tests/unit_tests/clients/agent_backend/test_factory.py b/api/tests/unit_tests/clients/agent_backend/test_factory.py index 0b595adda53..1c3a1878e98 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_factory.py +++ b/api/tests/unit_tests/clients/agent_backend/test_factory.py @@ -31,6 +31,7 @@ def test_create_agent_backend_client_forwards_authentication( client_cls.assert_called_once_with( base_url="http://agent-backend", stream_timeout=30, + binding_file_download_timeout=240, headers=headers, ) @@ -51,23 +52,33 @@ def test_create_agent_backend_run_client_forwards_stream_read_timeout(create_cli @pytest.mark.parametrize( - ("factory", "module"), + ("factory", "module", "extra_kwargs"), [ - (home_snapshot_service.AgentHomeSnapshotService._client, home_snapshot_service), - (workspace_service.AgentWorkspaceService._client, workspace_service), - (agent_app_sandbox_service._default_client_factory, agent_app_sandbox_service), + (home_snapshot_service.AgentHomeSnapshotService._client, home_snapshot_service, {}), + (workspace_service.AgentWorkspaceService._client, workspace_service, {}), + ( + agent_app_sandbox_service._default_client_factory, + agent_app_sandbox_service, + {"binding_file_download_timeout": 123.5}, + ), ], ) def test_default_agent_backend_clients_forward_authentication( monkeypatch: pytest.MonkeyPatch, factory: Callable[[], Client], module: ModuleType, + extra_kwargs: dict[str, float], ) -> None: monkeypatch.setattr(dify_config, "AGENT_BACKEND_BASE_URL", "http://agent-backend") monkeypatch.setattr(dify_config, "AGENT_BACKEND_API_TOKEN", "secret-token") + monkeypatch.setattr(dify_config, "AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS", 123.5) create_client = MagicMock() monkeypatch.setattr(module, "create_agent_backend_client", create_client) factory() - create_client.assert_called_once_with(base_url="http://agent-backend", api_token="secret-token") + create_client.assert_called_once_with( + base_url="http://agent-backend", + api_token="secret-token", + **extra_kwargs, + ) diff --git a/api/tests/unit_tests/commands/test_generate_swagger_specs.py b/api/tests/unit_tests/commands/test_generate_swagger_specs.py index 0ef5e7aca0d..00a8f157226 100644 --- a/api/tests/unit_tests/commands/test_generate_swagger_specs.py +++ b/api/tests/unit_tests/commands/test_generate_swagger_specs.py @@ -302,12 +302,11 @@ def test_generate_specs_include_console_contract_shapes_for_schema_migration(tmp assert app_list_item_schema["properties"]["can_trial"]["type"] == "boolean" assert "anyOf" not in app_list_item_schema["properties"]["can_trial"] assert "can_trial" in app_list_item_schema["required"] - app_detail_nullable_schema = schemas["RecommendedAppDetailNullableResponse"] assert _response_schema(paths["/explore/apps/{app_id}"]["get"])["$ref"] == ( - "#/components/schemas/RecommendedAppDetailNullableResponse" + "#/components/schemas/RecommendedAppDetailResponse" ) - assert {"$ref": "#/components/schemas/RecommendedAppDetailResponse"} in app_detail_nullable_schema["anyOf"] - assert {"type": "null"} in app_detail_nullable_schema["anyOf"] + assert "404" in paths["/explore/apps/{app_id}"]["get"]["responses"] + assert "RecommendedAppDetailNullableResponse" not in schemas assert schemas["RecommendedAppInfoResponse"]["properties"]["icon_url"]["readOnly"] is True assert schemas["InstalledAppInfoResponse"]["properties"]["icon_url"]["readOnly"] is True assert _response_schema(paths["/apps/{app_id}"]["get"])["$ref"] == "#/components/schemas/AppDetailWithSite" diff --git a/api/tests/unit_tests/configs/test_agent_backend_config.py b/api/tests/unit_tests/configs/test_agent_backend_config.py new file mode 100644 index 00000000000..542e2934ced --- /dev/null +++ b/api/tests/unit_tests/configs/test_agent_backend_config.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from configs.extra.agent_backend_config import AgentBackendConfig + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_API_TIMEOUT_ENV = "AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS" +_AGENT_TIMEOUT_ENV = "DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS" + + +def test_binding_file_download_timeout_defaults_to_240_seconds(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(_API_TIMEOUT_ENV, raising=False) + + assert AgentBackendConfig().AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS == 240.0 + + +def test_binding_file_download_timeout_rejects_non_positive_values() -> None: + with pytest.raises(ValidationError): + AgentBackendConfig(AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS=0) + + +def test_binding_file_timeout_docker_settings_use_their_service_env_files() -> None: + root_env_example = (_REPOSITORY_ROOT / "docker/.env.example").read_text(encoding="utf-8") + api_env_example = (_REPOSITORY_ROOT / "docker/envs/core-services/api.env.example").read_text(encoding="utf-8") + agent_env_example = (_REPOSITORY_ROOT / "docker/envs/core-services/dify-agent.env.example").read_text( + encoding="utf-8" + ) + compose_template = (_REPOSITORY_ROOT / "docker/docker-compose-template.yaml").read_text(encoding="utf-8") + + assert f"{_API_TIMEOUT_ENV}=" not in root_env_example + assert f"{_AGENT_TIMEOUT_ENV}=" not in root_env_example + assert f"{_API_TIMEOUT_ENV}=" in api_env_example + assert f"{_AGENT_TIMEOUT_ENV}=" not in api_env_example + assert f"{_API_TIMEOUT_ENV}=" not in agent_env_example + assert f"{_AGENT_TIMEOUT_ENV}=" in agent_env_example + assert f"{_API_TIMEOUT_ENV}:" not in compose_template + assert f"{_AGENT_TIMEOUT_ENV}:" not in compose_template diff --git a/api/tests/unit_tests/configs/test_logstore_config.py b/api/tests/unit_tests/configs/test_logstore_config.py new file mode 100644 index 00000000000..960e50b5a3a --- /dev/null +++ b/api/tests/unit_tests/configs/test_logstore_config.py @@ -0,0 +1,25 @@ +import pytest + +from configs.extra.logstore_config import LogStoreConfig +from tests.unit_tests.configs._isolated_settings import InitSettingsOnly + + +class _IsolatedLogStoreConfig(InitSettingsOnly, LogStoreConfig): + pass + + +@pytest.mark.parametrize( + ("raw_value", "expected"), + [ + pytest.param("true", True, id="enabled"), + pytest.param("false", False, id="disabled"), + ], +) +def test_logstore_migration_flags_parse_boolean_values(raw_value: str, expected: bool) -> None: + config = _IsolatedLogStoreConfig( + LOGSTORE_DUAL_WRITE_ENABLED=raw_value, + LOGSTORE_ENABLE_PUT_GRAPH_FIELD=raw_value, + ) + + assert config.LOGSTORE_DUAL_WRITE_ENABLED is expected + assert config.LOGSTORE_ENABLE_PUT_GRAPH_FIELD is expected diff --git a/api/tests/unit_tests/controllers/console/app/test_audio.py b/api/tests/unit_tests/controllers/console/app/test_audio.py index 347627809ff..e9f01b00c48 100644 --- a/api/tests/unit_tests/controllers/console/app/test_audio.py +++ b/api/tests/unit_tests/controllers/console/app/test_audio.py @@ -2,12 +2,12 @@ from __future__ import annotations import io from inspect import unwrap -from types import SimpleNamespace from unittest.mock import patch from uuid import UUID import pytest from flask import Flask +from sqlalchemy.orm import Session from werkzeug.datastructures import FileStorage from werkzeug.exceptions import Forbidden, InternalServerError @@ -34,7 +34,7 @@ from controllers.console.app.error import ( ) from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError -from models import AppMode +from models import Account, App, AppMode from models.agent import AgentConfigDraftType from models.agent_config_entities import AgentSoulConfig from services.agent.composer_service import AgentComposerService @@ -56,11 +56,30 @@ def _file_data(): return FileStorage(stream=io.BytesIO(b"audio"), filename="audio.wav", content_type="audio/wav") +def _app(*, app_id: str = "a1", tenant_id: str = "tenant-1") -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name="Audio app", + description="", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + + +def _account(account_id: str = "account-1") -> Account: + account = Account(name="Audio account", email=f"{account_id}@example.com") + account.id = account_id + return account + + def test_console_audio_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "ok"}) api = ChatMessageAudioApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") + app_model = _app() with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}): response = handler(api, app_model=app_model) @@ -72,9 +91,11 @@ def test_console_audio_api_accepts_published_agent_apps() -> None: assert AppMode.AGENT in audio_module._CONSOLE_AUDIO_TRANSCRIPT_APP_MODES -def test_agent_console_audio_api_uses_agent_draft(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_agent_console_audio_api_uses_agent_draft( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - app_model = SimpleNamespace(id="backing-app-1") + app_model = _app(app_id="backing-app-1") agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) calls: dict[str, object] = {} @@ -100,8 +121,8 @@ def test_agent_console_audio_api_uses_agent_draft(app: Flask, monkeypatch: pytes api = AgentChatMessageAudioApi() handler = unwrap(api.post) - session = SimpleNamespace() - current_user = SimpleNamespace(id="account-1") + session = unbound_session + current_user = _account() with app.test_request_context( f"/console/api/agent/{agent_id}/audio-to-text", method="POST", @@ -140,13 +161,15 @@ def test_agent_console_audio_api_uses_agent_draft(app: Flask, monkeypatch: pytes } -def test_agent_console_audio_api_defaults_to_normal_draft(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_agent_console_audio_api_defaults_to_normal_draft( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") captured: dict[str, object] = {} monkeypatch.setattr( audio_module, "resolve_agent_runtime_app_model", - lambda **_kwargs: SimpleNamespace(id="backing-app-1"), + lambda **_kwargs: _app(app_id="backing-app-1"), ) def load_agent_soul_for_debug(**kwargs): @@ -165,9 +188,9 @@ def test_agent_console_audio_api_defaults_to_normal_draft(app: Flask, monkeypatc ): response = handler( api, - session=SimpleNamespace(), + session=unbound_session, current_tenant_id="tenant-1", - current_user=SimpleNamespace(id="account-1"), + current_user=_account(), agent_id=agent_id, ) @@ -175,9 +198,11 @@ def test_agent_console_audio_api_defaults_to_normal_draft(app: Flask, monkeypatc assert captured["draft_type"] == AgentConfigDraftType.DRAFT -def test_agent_console_audio_api_checks_rbac_with_backing_app_id(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_agent_console_audio_api_checks_rbac_with_backing_app_id( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - app_model = SimpleNamespace(id="backing-app-1") + app_model = _app(app_id="backing-app-1") soul_loaded = False monkeypatch.setattr(audio_module, "resolve_agent_runtime_app_model", lambda **_kwargs: app_model) @@ -204,21 +229,23 @@ def test_agent_console_audio_api_checks_rbac_with_backing_app_id(app: Flask, mon with pytest.raises(Forbidden): handler( api, - session=SimpleNamespace(), + session=unbound_session, current_tenant_id="tenant-1", - current_user=SimpleNamespace(id="account-1"), + current_user=_account(), agent_id=agent_id, ) assert soul_loaded is False -def test_agent_console_audio_api_preserves_missing_build_draft_404(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_agent_console_audio_api_preserves_missing_build_draft_404( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") monkeypatch.setattr( audio_module, "resolve_agent_runtime_app_model", - lambda **_kwargs: SimpleNamespace(id="backing-app-1"), + lambda **_kwargs: _app(app_id="backing-app-1"), ) monkeypatch.setattr( AgentComposerService, @@ -236,9 +263,9 @@ def test_agent_console_audio_api_preserves_missing_build_draft_404(app: Flask, m with pytest.raises(AgentVersionNotFoundError): handler( api, - session=SimpleNamespace(), + session=unbound_session, current_tenant_id="tenant-1", - current_user=SimpleNamespace(id="account-1"), + current_user=_account(), agent_id=agent_id, ) @@ -262,7 +289,7 @@ def test_console_audio_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyP monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(exc)) api = ChatMessageAudioApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") + app_model = _app() with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}): with pytest.raises(expected): @@ -273,7 +300,7 @@ def test_console_audio_api_unhandled_error(app: Flask, monkeypatch: pytest.Monke monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("boom"))) api = ChatMessageAudioApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") + app_model = _app() with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}): with pytest.raises(InternalServerError): @@ -285,7 +312,7 @@ def test_console_text_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) - api = ChatMessageTextApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") + app_model = _app() with app.test_request_context( "/console/api/apps/app/text-to-audio", @@ -300,7 +327,7 @@ def test_console_text_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) - def test_console_text_api_builds_message_ref(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: api = ChatMessageTextApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") + app_model = _app(app_id="app-1") calls = {} def fake_transcript_tts(**kwargs): @@ -315,7 +342,7 @@ def test_console_text_api_builds_message_ref(app: Flask, monkeypatch: pytest.Mon method="POST", json={"text": "hello", "message_id": "message-1"}, ), - patch("controllers.console.app.audio.current_user", SimpleNamespace(id="account-1")), + patch("controllers.console.app.audio.current_user", _account()), ): response = handler(api, TextToSpeechPayload(text="hello", message_id="message-1"), app_model=app_model) @@ -328,7 +355,7 @@ def test_console_text_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyPa api = ChatMessageTextApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") + app_model = _app() with app.test_request_context( "/console/api/apps/app/text-to-audio", @@ -345,7 +372,7 @@ def test_console_text_modes_success(app: Flask, monkeypatch: pytest.MonkeyPatch) api = TextModesApi() handler = unwrap(api.get) - app_model = SimpleNamespace(tenant_id="t1") + app_model = _app(tenant_id="t1") with app.test_request_context("/console/api/apps/app/text-to-audio/voices?language=en", method="GET"): response = handler(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model) @@ -362,7 +389,7 @@ def test_console_text_modes_language_error(app: Flask, monkeypatch: pytest.Monke api = TextModesApi() handler = unwrap(api.get) - app_model = SimpleNamespace(tenant_id="t1") + app_model = _app(tenant_id="t1") with app.test_request_context("/console/api/apps/app/text-to-audio/voices?language=en", method="GET"): with pytest.raises(AppUnavailableError): @@ -376,7 +403,7 @@ def test_audio_to_text_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> N response_payload = {"text": "hello"} monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: response_payload) - app_model = SimpleNamespace(id="app-1") + app_model = _app(app_id="app-1") data = {"file": (io.BytesIO(b"x"), "sample.wav")} with app.test_request_context( @@ -400,7 +427,7 @@ def test_audio_to_text_maps_audio_too_large(app: Flask, monkeypatch: pytest.Monk lambda **_kwargs: (_ for _ in ()).throw(AudioTooLargeServiceError("too large")), ) - app_model = SimpleNamespace(id="app-1") + app_model = _app(app_id="app-1") data = {"file": (io.BytesIO(b"x"), "sample.wav")} with app.test_request_context( @@ -419,7 +446,7 @@ def test_text_to_audio_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> N monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "ok"}) - app_model = SimpleNamespace(id="app-1") + app_model = _app(app_id="app-1") with app.test_request_context( "/console/api/apps/app-1/text-to-audio", @@ -438,7 +465,7 @@ def test_text_to_audio_voices_success(app: Flask, monkeypatch: pytest.MonkeyPatc expected_voices = [{"name": "Voice 1", "value": "voice-1"}] monkeypatch.setattr(AudioService, "transcript_tts_voices", lambda **_kwargs: expected_voices) - app_model = SimpleNamespace(tenant_id="tenant-1") + app_model = _app() with app.test_request_context( "/console/api/apps/app-1/text-to-audio/voices", @@ -456,7 +483,7 @@ def test_audio_to_text_with_invalid_file(app: Flask, monkeypatch: pytest.MonkeyP monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "test"}) - app_model = SimpleNamespace(id="app-1") + app_model = _app(app_id="app-1") data = {"file": (io.BytesIO(b"invalid"), "sample.xyz")} with app.test_request_context( @@ -476,7 +503,7 @@ def test_text_to_audio_with_language_param(app: Flask, monkeypatch: pytest.Monke monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "test"}) - app_model = SimpleNamespace(id="app-1") + app_model = _app(app_id="app-1") with app.test_request_context( "/console/api/apps/app-1/text-to-audio", @@ -497,7 +524,7 @@ def test_text_to_audio_voices_with_language_filter(app: Flask, monkeypatch: pyte lambda **_kwargs: [{"name": "Voice 1", "value": "voice-1"}], ) - app_model = SimpleNamespace(tenant_id="tenant-1") + app_model = _app() with app.test_request_context( "/console/api/apps/app-1/text-to-audio/voices?language=en-US", diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow.py b/api/tests/unit_tests/controllers/console/app/test_workflow.py index 3677b0cbb80..bf59f8a32ad 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow.py @@ -2,6 +2,7 @@ from __future__ import annotations import inspect import json +from contextlib import contextmanager, nullcontext from datetime import datetime from types import SimpleNamespace from typing import cast @@ -79,6 +80,99 @@ def _make_workflow(**overrides): return workflow +def test_publish_workflow_returns_success( + app: Flask, + monkeypatch: pytest.MonkeyPatch, +) -> None: + current_user = SimpleNamespace(id="account-1") + app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") + workflow = SimpleNamespace(id="published-workflow", created_at=datetime(2026, 8, 17, 12, 0, 0)) + session = Mock() + session.get.return_value = app_model + monkeypatch.setattr( + workflow_module, + "WorkflowService", + Mock(return_value=SimpleNamespace(publish_workflow=Mock(return_value=workflow))), + ) + monkeypatch.setattr( + workflow_module, + "sessionmaker", + lambda _engine: SimpleNamespace(begin=lambda: nullcontext(session)), + ) + monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object())) + with app.test_request_context("/apps/app-1/workflows/publish", method="POST", json={}): + response = inspect.unwrap(workflow_module.PublishedWorkflowApi.post)( + workflow_module.PublishedWorkflowApi(), + current_user, + app_model, + ) + + assert response["result"] == "success" + + +@pytest.mark.parametrize("transaction_fails", [False, True], ids=["commit-succeeds", "commit-fails"]) +def test_delete_workflow_retires_candidates_only_after_transaction_exit( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + transaction_fails: bool, +) -> None: + current_user = SimpleNamespace(id="account-1") + app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") + session = Mock() + events: list[str] = [] + error = RuntimeError("commit failed") + workflow_service = SimpleNamespace( + delete_workflow=Mock(side_effect=lambda **_kwargs: events.append("delete") or ["inline-agent"]) + ) + + @contextmanager + def transaction(): + events.append("transaction-enter") + yield session + events.append("transaction-exit") + if transaction_fails: + raise error + + monkeypatch.setattr(workflow_module, "WorkflowService", Mock(return_value=workflow_service)) + monkeypatch.setattr( + workflow_module, + "sessionmaker", + lambda _engine: SimpleNamespace(begin=transaction), + ) + monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object())) + retire_unowned = Mock(side_effect=lambda **_kwargs: events.append("retire")) + monkeypatch.setattr(workflow_module.WorkflowAgentRetirementService, "retire_unowned", retire_unowned) + + with app.test_request_context("/apps/app-1/workflows/workflow-1", method="DELETE"): + if transaction_fails: + with pytest.raises(RuntimeError) as exc_info: + inspect.unwrap(workflow_module.WorkflowByIdApi.delete)( + workflow_module.WorkflowByIdApi(), + current_user, + app_model, + "workflow-1", + ) + assert exc_info.value is error + else: + response = inspect.unwrap(workflow_module.WorkflowByIdApi.delete)( + workflow_module.WorkflowByIdApi(), + current_user, + app_model, + "workflow-1", + ) + assert response == (None, 204) + + assert events == ["transaction-enter", "delete", "transaction-exit"] + ([] if transaction_fails else ["retire"]) + if transaction_fails: + retire_unowned.assert_not_called() + else: + retire_unowned.assert_called_once_with( + tenant_id=app_model.tenant_id, + agent_ids=["inline-agent"], + account_id=current_user.id, + ) + + def test_parse_file_no_config(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(workflow_module.FileUploadConfigManager, "convert", lambda *_args, **_kwargs: None) workflow = SimpleNamespace(features_dict={}, tenant_id="t1") @@ -752,12 +846,19 @@ def test_workflow_online_users_filters_inaccessible_workflow(app: Flask, monkeyp app_id_2 = "22222222-2222-2222-2222-222222222222" signed_avatar_url = "https://files.example.com/signed/avatar-1" sign_avatar = Mock(return_value=signed_avatar_url) + get_tenant_app_maintainers = Mock(return_value={app_id_1: "owner-1", app_id_2: "owner-2"}) monkeypatch.setattr( workflow_module, "WorkflowService", - lambda: SimpleNamespace(get_accessible_app_ids=lambda app_ids, tenant_id, session: {app_id_1}), + lambda: SimpleNamespace(get_tenant_app_maintainers=get_tenant_app_maintainers), ) + access_filter = SimpleNamespace(is_app_accessible=lambda app_id, _maintainer, _account_id: app_id == app_id_1) + resolve_access = Mock(return_value=access_filter) + monkeypatch.setattr(workflow_module, "resolve_app_access_filter", resolve_access) + monkeypatch.setattr(workflow_module.dify_config, "RBAC_ENABLED", True) monkeypatch.setattr(workflow_module.file_helpers, "get_signed_file_url", sign_avatar) + short_session = Mock() + monkeypatch.setattr(workflow_module.session_factory, "create_session", lambda: nullcontext(short_session)) redis_pipeline = Mock() redis_pipeline.execute.return_value = [ @@ -805,7 +906,7 @@ def test_workflow_online_users_filters_inaccessible_workflow(app: Flask, monkeyp method="POST", json={"app_ids": [app_id_1, app_id_2]}, ): - response = handler(api, "tenant-1") + response = handler(api, "tenant-1", SimpleNamespace(id="account-1")) assert response == { "data": [ @@ -830,6 +931,11 @@ def test_workflow_online_users_filters_inaccessible_workflow(app: Flask, monkeyp redis_pipeline.hgetall.assert_called_once_with(f"{workflow_module.WORKFLOW_ONLINE_USERS_PREFIX}{app_id_1}") redis_pipeline.execute.assert_called_once_with() sign_avatar.assert_called_once_with("avatar-file-id") + get_tenant_app_maintainers.assert_called_once() + resolve_access.assert_called_once() + assert get_tenant_app_maintainers.call_args.args == ([app_id_1, app_id_2], "tenant-1") + assert resolve_access.call_args.args == ("tenant-1", "account-1") + assert get_tenant_app_maintainers.call_args.kwargs["session"] is resolve_access.call_args.kwargs["session"] def test_workflow_online_users_batches_redis_reads(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: @@ -837,8 +943,10 @@ def test_workflow_online_users_batches_redis_reads(app: Flask, monkeypatch: pyte monkeypatch.setattr( workflow_module, "WorkflowService", - lambda: SimpleNamespace(get_accessible_app_ids=lambda app_ids, tenant_id, session: set(app_ids)), + lambda: SimpleNamespace(get_tenant_app_maintainers=lambda app_ids, tenant_id, session: dict.fromkeys(app_ids)), ) + monkeypatch.setattr(workflow_module.dify_config, "RBAC_ENABLED", False) + monkeypatch.setattr(workflow_module.session_factory, "create_session", lambda: nullcontext(Mock())) first_pipeline = Mock() first_pipeline.execute.return_value = [{} for _ in range(workflow_module.WORKFLOW_ONLINE_USERS_REDIS_BATCH_SIZE)] @@ -855,7 +963,7 @@ def test_workflow_online_users_batches_redis_reads(app: Flask, monkeypatch: pyte method="POST", json={"app_ids": app_ids}, ): - response = handler(api, "tenant-1") + response = handler(api, "tenant-1", SimpleNamespace(id="account-1")) assert len(response["data"]) == len(app_ids) assert redis_pipeline_factory.call_count == 2 @@ -864,11 +972,11 @@ def test_workflow_online_users_batches_redis_reads(app: Flask, monkeypatch: pyte def test_workflow_online_users_rejects_excessive_workflow_ids(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - accessible_app_ids = Mock(return_value=set()) + get_tenant_app_maintainers = Mock(return_value={}) monkeypatch.setattr( workflow_module, "WorkflowService", - lambda: SimpleNamespace(get_accessible_app_ids=accessible_app_ids), + lambda: SimpleNamespace(get_tenant_app_maintainers=get_tenant_app_maintainers), ) excessive_ids = [f"wf-{index}" for index in range(workflow_module.MAX_WORKFLOW_ONLINE_USERS_REQUEST_IDS + 1)] @@ -882,9 +990,9 @@ def test_workflow_online_users_rejects_excessive_workflow_ids(app: Flask, monkey json={"app_ids": excessive_ids}, ): with pytest.raises(HTTPException) as exc: - handler(api, "tenant-1") + handler(api, "tenant-1", SimpleNamespace(id="account-1")) assert exc.value.code == 400 assert exc.value.description is not None assert "Maximum" in exc.value.description - accessible_app_ids.assert_not_called() + get_tenant_app_maintainers.assert_not_called() diff --git a/api/tests/unit_tests/controllers/console/app/test_wraps.py b/api/tests/unit_tests/controllers/console/app/test_wraps.py index 2f94aedaf52..e341dc15d84 100644 --- a/api/tests/unit_tests/controllers/console/app/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/app/test_wraps.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock from uuid import uuid4 import pytest +from sqlalchemy import event, text from sqlalchemy.orm import Session from controllers.common import session as session_module @@ -15,7 +16,7 @@ from controllers.console.app import completion as completion_module from controllers.console.app import workflow as workflow_module from controllers.console.app import wraps as wraps_module from controllers.console.app.error import AppNotFoundError -from models.model import App, AppMode, TrialApp +from models.model import App, AppMode def _persist_app(sqlite_session: Session, *, mode: AppMode = AppMode.CHAT) -> App: @@ -57,60 +58,58 @@ def test_get_app_model_rejects_wrong_mode(monkeypatch: pytest.MonkeyPatch, sqlit handler(app_id=app_model.id) -def test_get_app_model_with_trial_requires_trial_app_registration( +def test_load_previewable_app_model_rejects_app_outside_preview_admission( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = MagicMock(spec=Session) + app_loader = MagicMock() + recommended_app_queries = MagicMock() + recommended_app_queries.is_previewable.return_value = False + monkeypatch.setattr( + wraps_module, + "application_services", + lambda: SimpleNamespace(recommended_app_queries=recommended_app_queries), + ) + monkeypatch.setattr(wraps_module.AppService, "get_normal_app_by_id", app_loader) + + assert wraps_module._load_previewable_app_model(session, "app-1") is None + recommended_app_queries.is_previewable.assert_called_once_with("app-1") + app_loader.assert_not_called() + + +def test_load_previewable_app_model_rejects_non_normal_app( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: app_model = _persist_app(sqlite_session) - recommended_get_app = MagicMock(return_value=None) - monkeypatch.setattr(wraps_module.RecommendedAppService, "get_app", recommended_get_app) + app_id = app_model.id + sqlite_session.execute(text("UPDATE apps SET status = 'disabled' WHERE id = :app_id"), {"app_id": app_id}) + sqlite_session.commit() + recommended_app_queries = MagicMock() + recommended_app_queries.is_previewable.return_value = True + monkeypatch.setattr( + wraps_module, + "application_services", + lambda: SimpleNamespace(recommended_app_queries=recommended_app_queries), + ) + + assert wraps_module._load_previewable_app_model(sqlite_session, app_id) is None + + +def test_get_previewable_app_model_rejects_app_outside_preview_admission( + monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: + app_loader = MagicMock(return_value=None) + monkeypatch.setattr(wraps_module, "_load_previewable_app_model", app_loader) class Handler: - @wraps_module.get_app_model_with_trial + @wraps_module.get_previewable_app_model def get(self, _injected_session, app_model): return app_model.id with pytest.raises(AppNotFoundError): - Handler().get(sqlite_session, app_id=app_model.id) + Handler().get(unbound_session, app_id="app-1") - recommended_get_app.assert_called_once_with(app_model.id, session=sqlite_session) - - -def test_get_app_model_with_trial_falls_back_to_recommended_app( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - app_model = SimpleNamespace(id="app-1", mode=AppMode.CHAT.value, status="normal", tenant_id="t1") - trial_app_loader = MagicMock(return_value=None) - recommended_get_app = MagicMock(return_value=app_model) - monkeypatch.setattr(wraps_module, "_load_app_model_with_trial", trial_app_loader) - monkeypatch.setattr(wraps_module.RecommendedAppService, "get_app", recommended_get_app) - - class Handler: - @wraps_module.get_app_model_with_trial - def get(self, _injected_session, app_model): - return app_model.id - - assert Handler().get(unbound_session, app_id="app-1") == "app-1" - trial_app_loader.assert_called_once_with(unbound_session, "app-1") - recommended_get_app.assert_called_once_with("app-1", session=unbound_session) - - -def test_get_app_model_with_trial_prefers_trial_registration( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - app_model = SimpleNamespace(id="app-1", mode=AppMode.CHAT.value, status="normal", tenant_id="t1") - trial_app_loader = MagicMock(return_value=app_model) - recommended_get_app = MagicMock() - monkeypatch.setattr(wraps_module, "_load_app_model_with_trial", trial_app_loader) - monkeypatch.setattr(wraps_module.RecommendedAppService, "get_app", recommended_get_app) - - class Handler: - @wraps_module.get_app_model_with_trial - def get(self, _injected_session, app_model): - return app_model.id - - assert Handler().get(unbound_session, app_id="app-1") == "app-1" - trial_app_loader.assert_called_once_with(unbound_session, "app-1") - recommended_get_app.assert_not_called() + app_loader.assert_called_once_with(unbound_session, "app-1") def test_get_app_model_requires_app_id() -> None: @@ -145,12 +144,32 @@ def test_get_app_model_prefers_injected_session( assert Handler().get(sqlite_session, app_id=app_model.id) == app_model.id -def test_get_app_model_with_trial_prefers_injected_session( +def test_preview_admission_precedes_request_session_transaction( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: app_model = _persist_app(sqlite_session) - sqlite_session.add(TrialApp(app_id=app_model.id, tenant_id=app_model.tenant_id)) - sqlite_session.commit() + app_id = app_model.id + sqlite_session.rollback() + request_transaction_begins = 0 + + def record_request_transaction_begin(_session, _transaction, _connection) -> None: + nonlocal request_transaction_begins + request_transaction_begins += 1 + + event.listen(sqlite_session, "after_begin", record_request_transaction_begin) + recommended_app_queries = MagicMock() + + def assert_request_session_has_not_started(_app_id: str) -> bool: + assert request_transaction_begins == 0 + assert sqlite_session.in_transaction() is False + return True + + recommended_app_queries.is_previewable.side_effect = assert_request_session_has_not_started + monkeypatch.setattr( + wraps_module, + "application_services", + lambda: SimpleNamespace(recommended_app_queries=recommended_app_queries), + ) monkeypatch.setattr( wraps_module.db, "session", @@ -160,16 +179,18 @@ def test_get_app_model_with_trial_prefers_injected_session( class Handler: @with_session(write=False) - @wraps_module.get_app_model_with_trial(None) + @wraps_module.get_previewable_app_model(None) def get(self, injected_session, app_model): assert injected_session is sqlite_session return app_model.id - assert Handler().get(app_id=app_model.id) == app_model.id + assert Handler().get(app_id=app_id) == app_id + recommended_app_queries.is_previewable.assert_called_once_with(app_id) + assert request_transaction_begins == 1 -def test_get_app_model_with_trial_requires_injected_session() -> None: - @wraps_module.get_app_model_with_trial(None) +def test_get_previewable_app_model_requires_injected_session() -> None: + @wraps_module.get_previewable_app_model(None) def handler(app_model): return app_model.id diff --git a/api/tests/unit_tests/controllers/console/auth/test_account_activation.py b/api/tests/unit_tests/controllers/console/auth/test_account_activation.py index bd5bd18e089..6f5ba50e56a 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_account_activation.py +++ b/api/tests/unit_tests/controllers/console/auth/test_account_activation.py @@ -9,13 +9,22 @@ from flask import Flask from controllers.console.auth.activate import ActivateApi, ActivateCheckApi from controllers.console.auth.error import InvitationAccountMismatchError as InvitationAccountMismatchHTTPError -from controllers.console.error import AccountInFreezeError, AlreadyActivateError +from controllers.console.error import ( + AccountInFreezeError, + AlreadyActivateError, +) +from controllers.console.error import ( + EmailDomainSuspendedError as EmailDomainSuspendedHTTPError, +) from services.account_activation_service import ( AccountActivationService, FrozenAccountError, InvalidInvitationError, InvitationAccountMismatchError, ) +from services.account_activation_service import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) from services.entities.account_activation_entities import ( ActivationCheckData, ActivationCheckResult, @@ -172,6 +181,7 @@ class TestActivateApi: (InvalidInvitationError(), AlreadyActivateError), (InvitationAccountMismatchError(), InvitationAccountMismatchHTTPError), (FrozenAccountError(), AccountInFreezeError), + (EmailDomainSuspendedRegistrationError(), EmailDomainSuspendedHTTPError), ], ) def test_translates_application_errors( diff --git a/api/tests/unit_tests/controllers/console/auth/test_email_register.py b/api/tests/unit_tests/controllers/console/auth/test_email_register.py index 9a8b4db80d0..74a2fbaf321 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_email_register.py +++ b/api/tests/unit_tests/controllers/console/auth/test_email_register.py @@ -4,6 +4,7 @@ from __future__ import annotations from unittest.mock import MagicMock, patch +import pytest from flask import Flask from controllers.console.auth.email_register import ( @@ -11,14 +12,21 @@ from controllers.console.auth.email_register import ( EmailRegisterResetApi, EmailRegisterSendEmailApi, ) +from controllers.console.error import AccountInFreezeError, EmailDomainSuspendedError from enums import DeploymentEdition from services.entities.feature_entities import SystemFeatureModel +from services.errors.account import ( + AccountRegisterError, +) +from services.errors.account import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) class TestEmailRegisterSendEmailApi: @patch("controllers.console.auth.email_register.AccountService.get_account_by_email_with_case_fallback") @patch("controllers.console.auth.email_register.AccountService.send_email_register_email") - @patch("controllers.console.auth.email_register.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.email_register.BillingService.get_email_freeze_type") @patch("controllers.console.auth.email_register.AccountService.is_email_send_ip_limit", return_value=False) @patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1") def test_send_email_normalizes_and_falls_back( @@ -31,7 +39,7 @@ class TestEmailRegisterSendEmailApi: app: Flask, ): mock_send_mail.return_value = "token-123" - mock_is_freeze.return_value = False + mock_is_freeze.return_value = None mock_account = MagicMock() mock_get_account.return_value = mock_account @@ -58,6 +66,49 @@ class TestEmailRegisterSendEmailApi: mock_extract_ip.assert_called_once() mock_is_email_send_ip_limit.assert_called_once_with("127.0.0.1") + @pytest.mark.parametrize( + ("freeze_type", "expected_error"), + [ + ("freeze", AccountInFreezeError), + ("email_domain_suspended", EmailDomainSuspendedError), + ], + ) + @patch("controllers.console.auth.email_register.BillingService.get_email_freeze_type") + @patch("controllers.console.auth.email_register.AccountService.is_email_send_ip_limit", return_value=False) + @patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1") + def test_send_email_rejects_frozen_email( + self, + mock_extract_ip, + mock_is_email_send_ip_limit, + mock_get_freeze_type, + app: Flask, + freeze_type, + expected_error, + ): + mock_get_freeze_type.return_value = freeze_type + feature_flags = SystemFeatureModel( + deployment_edition=DeploymentEdition.COMMUNITY, + enable_email_password_login=True, + is_allow_register=True, + ) + + with ( + patch("controllers.console.auth.email_register.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD), + patch("controllers.console.wraps.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD), + patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags), + ): + with app.test_request_context( + "/email-register/send-email", + method="POST", + json={"email": "Invitee@Example.com"}, + ): + with pytest.raises(expected_error): + EmailRegisterSendEmailApi().post() + + mock_get_freeze_type.assert_called_once_with("invitee@example.com") + mock_is_email_send_ip_limit.assert_called_once_with("127.0.0.1") + mock_extract_ip.assert_called_once() + class TestEmailRegisterCheckApi: @patch("controllers.console.auth.email_register.AccountService.reset_email_register_error_rate_limit") @@ -107,6 +158,28 @@ class TestEmailRegisterCheckApi: class TestEmailRegisterResetApi: + @pytest.mark.parametrize( + ("service_error", "expected_error"), + [ + (EmailDomainSuspendedRegistrationError(), EmailDomainSuspendedError), + (AccountRegisterError("frozen"), AccountInFreezeError), + ], + ) + @patch("controllers.console.auth.email_register.AccountService.create_account_and_tenant") + def test_create_new_account_translates_freeze_errors( + self, + mock_create_account, + service_error, + expected_error, + ): + mock_create_account.side_effect = service_error + + with pytest.raises(expected_error): + EmailRegisterResetApi()._create_new_account( + email="user@example.com", + password="ValidPass123!", + ) + @patch("controllers.console.auth.email_register.AccountService.reset_login_error_rate_limit") @patch("controllers.console.auth.email_register.AccountService.login") @patch("controllers.console.auth.email_register.EmailRegisterResetApi._create_new_account") diff --git a/api/tests/unit_tests/controllers/console/auth/test_email_verification.py b/api/tests/unit_tests/controllers/console/auth/test_email_verification.py index 573934239f5..5fa6e48e62a 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_email_verification.py +++ b/api/tests/unit_tests/controllers/console/auth/test_email_verification.py @@ -33,6 +33,7 @@ from controllers.console.auth.login import ( from controllers.console.error import ( AccountInFreezeError, AccountNotFound, + EmailDomainSuspendedError, EmailSendIpLimitError, NotAllowedCreateWorkspace, WorkspacesLimitExceeded, @@ -43,7 +44,12 @@ from services.email_code_login_challenge import ( EmailCodeLoginChallengeStatus, EmailCodeLoginChallengeUnavailableError, ) -from services.errors.account import AccountRegisterError +from services.errors.account import ( + AccountRegisterError, +) +from services.errors.account import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) from services.turnstile_service import TurnstileChallengeRejectedError, TurnstileUpstreamError TEST_TOKEN = "00000000-0000-4000-8000-000000000001" @@ -308,7 +314,22 @@ class TestEmailCodeLoginSendEmailApi: @patch("controllers.console.wraps.db") @patch("controllers.console.auth.login.AccountService.is_email_send_ip_limit") @patch("controllers.console.auth.login.AccountService.get_user_through_email") - def test_send_email_code_frozen_account(self, mock_get_user, mock_is_ip_limit, mock_db, app: Flask): + @pytest.mark.parametrize( + ("service_error", "expected_error"), + [ + (AccountRegisterError("Account frozen"), AccountInFreezeError), + (EmailDomainSuspendedRegistrationError(), EmailDomainSuspendedError), + ], + ) + def test_send_email_code_frozen_account( + self, + mock_get_user, + mock_is_ip_limit, + mock_db, + app: Flask, + service_error, + expected_error, + ): """ Test email code sending to frozen account. @@ -317,12 +338,12 @@ class TestEmailCodeLoginSendEmailApi: """ # Arrange mock_is_ip_limit.return_value = False - mock_get_user.side_effect = AccountRegisterError("Account frozen") + mock_get_user.side_effect = service_error # Act & Assert with app.test_request_context("/email-code-login", method="POST", json={"email": "frozen@example.com"}): api = EmailCodeLoginSendEmailApi() - with pytest.raises(AccountInFreezeError): + with pytest.raises(expected_error): api.post() @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/controllers/console/auth/test_login_logout.py b/api/tests/unit_tests/controllers/console/auth/test_login_logout.py index a4da84d1c77..f6c0819c9d3 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_login_logout.py +++ b/api/tests/unit_tests/controllers/console/auth/test_login_logout.py @@ -22,17 +22,26 @@ from controllers.console.auth.error import ( EmailPasswordLoginLimitError, InvalidEmailError, ) -from controllers.console.auth.login import EmailCodeLoginApi, LoginApi, LogoutApi +from controllers.console.auth.login import EmailCodeLoginApi, LoginApi, LogoutApi, ResetPasswordSendEmailApi from controllers.console.error import ( AccountBannedError, AccountInFreezeError, + EmailDomainSuspendedError, SeatsLimitExceeded, WorkspacesLimitExceeded, ) from enums import DeploymentEdition from services.email_code_login_challenge import EmailCodeLoginChallengeResult, EmailCodeLoginChallengeStatus from services.entities.auth_entities import LoginFailureReason -from services.errors.account import AccountLoginError, AccountPasswordError, SeatsLimitExceededError +from services.errors.account import ( + AccountLoginError, + AccountPasswordError, + AccountRegisterError, + SeatsLimitExceededError, +) +from services.errors.account import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) TEST_TOKEN = "00000000-0000-4000-8000-000000000001" @@ -228,7 +237,7 @@ class TestLoginApi: @patch("controllers.console.wraps.db") @patch("controllers.console.auth.login.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD) - @patch("controllers.console.auth.login.BillingService.is_email_in_freeze") + @patch("controllers.console.auth.login.BillingService.get_email_freeze_type") def test_login_fails_when_account_frozen( self, mock_is_frozen, mock_db, app: Flask, caplog: pytest.LogCaptureFixture ): @@ -240,7 +249,7 @@ class TestLoginApi: - AccountInFreezeError is raised for frozen accounts """ # Arrange - mock_is_frozen.return_value = True + mock_is_frozen.return_value = "freeze" # Act & Assert with app.test_request_context( @@ -257,6 +266,116 @@ class TestLoginApi: assert warn_records[0].args[0] == "frozen@example.com" assert warn_records[0].args[1] == LoginFailureReason.ACCOUNT_IN_FREEZE + @patch("controllers.console.wraps.db") + @patch("controllers.console.auth.login.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD) + @patch("controllers.console.auth.login.BillingService.get_email_freeze_type") + def test_login_fails_when_email_domain_is_suspended(self, mock_get_freeze_type, mock_db, app: Flask): + mock_get_freeze_type.return_value = "email_domain_suspended" + + with app.test_request_context( + "/login", + method="POST", + json={"email": "user@suspended.example", "password": encode_password("password")}, + ): + with pytest.raises(EmailDomainSuspendedError): + LoginApi().post() + + @pytest.mark.parametrize( + ("service_error", "expected_error"), + [ + (EmailDomainSuspendedRegistrationError(), EmailDomainSuspendedError), + (AccountRegisterError("frozen"), AccountInFreezeError), + ], + ) + @patch("controllers.console.wraps.db") + @patch("controllers.console.auth.login._get_account_with_case_fallback") + @patch("controllers.console.auth.login.AccountService.verify_email_code_login_challenge") + def test_email_code_login_translates_freeze_errors( + self, + mock_verify_challenge, + mock_get_account, + mock_db, + app: Flask, + service_error, + expected_error, + ): + mock_verify_challenge.return_value = EmailCodeLoginChallengeResult( + status=EmailCodeLoginChallengeStatus.VERIFIED + ) + mock_get_account.side_effect = service_error + + with app.test_request_context( + "/email-code-login/validity", + method="POST", + json={"email": "User@Example.com", "code": encode_code("123456"), "token": TEST_TOKEN}, + ): + with pytest.raises(expected_error): + EmailCodeLoginApi().post() + + @pytest.mark.parametrize( + ("service_error", "expected_error"), + [ + (EmailDomainSuspendedRegistrationError(), EmailDomainSuspendedError), + (AccountRegisterError("frozen"), AccountInFreezeError), + ], + ) + @patch("controllers.console.wraps.db") + @patch("controllers.console.auth.login.db") + @patch("controllers.console.auth.login.AccountService.create_account_and_tenant") + @patch("controllers.console.auth.login.AccountService.verify_email_code_login_challenge") + @patch("controllers.console.auth.login._get_account_with_case_fallback") + def test_email_code_login_translates_account_creation_freeze_errors( + self, + mock_get_account, + mock_verify_challenge, + mock_create_account, + mock_login_db, + mock_db, + app: Flask, + service_error, + expected_error, + ): + mock_verify_challenge.return_value = EmailCodeLoginChallengeResult( + status=EmailCodeLoginChallengeStatus.VERIFIED + ) + mock_get_account.return_value = None + mock_create_account.side_effect = service_error + + with app.test_request_context( + "/email-code-login/validity", + method="POST", + json={"email": "User@Example.com", "code": encode_code("123456"), "token": TEST_TOKEN}, + ): + with pytest.raises(expected_error): + EmailCodeLoginApi().post() + + @pytest.mark.parametrize( + ("service_error", "expected_error"), + [ + (EmailDomainSuspendedRegistrationError(), EmailDomainSuspendedError), + (AccountRegisterError("frozen"), AccountInFreezeError), + ], + ) + @patch("controllers.console.wraps.db") + @patch("controllers.console.auth.login._get_account_with_case_fallback") + def test_reset_password_translates_freeze_errors( + self, + mock_get_account, + mock_db, + app: Flask, + service_error, + expected_error, + ): + mock_get_account.side_effect = service_error + + with app.test_request_context( + "/reset-password", + method="POST", + json={"email": "User@Example.com"}, + ): + with pytest.raises(expected_error): + ResetPasswordSendEmailApi().post() + @patch("controllers.console.wraps.db") @patch("controllers.console.auth.login.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY) @patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit") diff --git a/api/tests/unit_tests/controllers/console/auth/test_oauth.py b/api/tests/unit_tests/controllers/console/auth/test_oauth.py index 05fa93aa3b3..98d03ee2ac7 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_oauth.py +++ b/api/tests/unit_tests/controllers/console/auth/test_oauth.py @@ -14,9 +14,15 @@ from controllers.console.auth.oauth import ( _get_account_by_openid_or_email, get_oauth_providers, ) +from enums import DeploymentEdition from libs.oauth import OAuthUserInfo, encode_oauth_state from models.account import AccountStatus -from services.errors.account import AccountRegisterError +from services.errors.account import ( + AccountRegisterError, +) +from services.errors.account import ( + EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, +) @pytest.fixture(autouse=True) @@ -231,6 +237,38 @@ class TestOAuthCallback: ) mock_redirect.assert_called_once_with("http://localhost:3000?oauth_new_user=true") + @pytest.mark.parametrize( + ("service_error", "expected_message"), + [ + ( + EmailDomainSuspendedRegistrationError(), + "This email domain has been suspended.", + ), + (AccountRegisterError("This email account is frozen."), "This email account is frozen."), + ], + ) + @patch("controllers.console.auth.oauth.get_oauth_providers") + @patch("controllers.console.auth.oauth._generate_account") + @patch("controllers.console.auth.oauth.redirect") + def test_should_translate_registration_freeze_errors( + self, + mock_redirect, + mock_generate_account, + mock_get_providers, + resource: OAuthCallback, + app: Flask, + oauth_setup, + service_error, + expected_message, + ): + mock_get_providers.return_value = {"github": oauth_setup["provider"]} + mock_generate_account.side_effect = service_error + + with app.test_request_context("/auth/oauth/github/callback?code=test_code"): + resource.get("github") + + mock_redirect.assert_called_once_with(f"http://localhost:3000/signin?message={expected_message}") + @pytest.mark.parametrize( ("exception", "expected_error"), [ @@ -537,6 +575,36 @@ class TestAccountGeneration: else: mock_register_service.register.assert_not_called() + @pytest.mark.parametrize( + ("freeze_type", "expected_error"), + [ + ("email_domain_suspended", EmailDomainSuspendedRegistrationError), + ("freeze", AccountRegisterError), + ], + ) + @patch("controllers.console.auth.oauth.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD) + @patch("controllers.console.auth.oauth.BillingService.get_email_freeze_type") + @patch("controllers.console.auth.oauth._get_account_by_openid_or_email", return_value=None) + @patch("controllers.console.auth.oauth.FeatureService") + def test_should_reject_registration_for_frozen_email( + self, + mock_feature_service, + mock_get_account, + mock_get_freeze_type, + freeze_type, + expected_error, + app: Flask, + user_info: OAuthUserInfo, + ): + mock_feature_service.get_system_features.return_value.is_allow_register = False + mock_get_freeze_type.return_value = freeze_type + + with app.test_request_context("/"): + with pytest.raises(expected_error): + _generate_account("github", user_info) + + mock_get_freeze_type.assert_called_once_with("test@example.com") + @patch("controllers.console.auth.oauth._get_account_by_openid_or_email", return_value=None) @patch("controllers.console.auth.oauth.FeatureService") @patch("controllers.console.auth.oauth.RegisterService") diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py index caef20a3d69..81b58c32339 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py @@ -3,7 +3,8 @@ from __future__ import annotations import json -from collections.abc import Iterator +from collections.abc import Generator, Iterator +from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime from inspect import unwrap @@ -671,13 +672,33 @@ class TestRagPipelineByIdApi: result, status = method(api, WorkflowUpdatePayload(), user, pipeline, "w1") assert status == 400 - def test_delete_success(self, app: Flask) -> None: + @pytest.mark.parametrize("transaction_fails", [False, True], ids=["commit-succeeds", "commit-fails"]) + def test_delete_retires_candidates_only_after_transaction_exit( + self, + app: Flask, + transaction_fails: bool, + ) -> None: api = RagPipelineByIdApi() method = unwrap(api.delete) pipeline = make_pipeline(tenant_id="t1", workflow_id="active-workflow") + user = make_account() + events: list[str] = [] + error = RuntimeError("commit failed") workflow_service = MagicMock() + workflow_service.delete_workflow.side_effect = lambda **_kwargs: events.append("delete") or ["inline-agent"] + transaction_factory = MagicMock() + + @contextmanager + def transaction() -> Generator[Session]: + events.append("transaction-enter") + yield MagicMock(spec=Session) + events.append("transaction-exit") + if transaction_fails: + raise error + + transaction_factory.begin.side_effect = transaction with ( app.test_request_context("/", method="DELETE"), @@ -685,21 +706,45 @@ class TestRagPipelineByIdApi: "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.WorkflowService", return_value=workflow_service, ), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.sessionmaker", + return_value=transaction_factory, + ), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow." + "WorkflowAgentRetirementService.retire_unowned" + ) as retire_unowned, ): - result = method(api, pipeline, "old-workflow") + retire_unowned.side_effect = lambda **_kwargs: events.append("retire") + if transaction_fails: + with pytest.raises(RuntimeError) as exc_info: + method(api, user, pipeline, "old-workflow") + assert exc_info.value is error + else: + result = method(api, user, pipeline, "old-workflow") + assert result == (None, 204) workflow_service.delete_workflow.assert_called_once() - assert result == (None, 204) + assert events == ["transaction-enter", "delete", "transaction-exit"] + ([] if transaction_fails else ["retire"]) + if transaction_fails: + retire_unowned.assert_not_called() + else: + retire_unowned.assert_called_once_with( + tenant_id=pipeline.tenant_id, + agent_ids=["inline-agent"], + account_id=user.id, + ) def test_delete_active_workflow_rejected(self, app: Flask) -> None: api = RagPipelineByIdApi() method = unwrap(api.delete) pipeline = make_pipeline(tenant_id="t1", workflow_id="active-workflow") + user = make_account() with app.test_request_context("/", method="DELETE"): with pytest.raises(BadRequest, match="currently in use by pipeline"): - method(api, pipeline, "active-workflow") + method(api, user, pipeline, "active-workflow") class TestRagPipelineWorkflowLastRunApi: diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py index bbcda2fb372..b2485d55d0f 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py @@ -587,17 +587,20 @@ class TestDatasetDocumentSegmentBatchImportApi: used=False, ) user = MagicMock(id="u1") + dataset = MagicMock(id="ds-1", tenant_id="tenant-1") session = MagicMock() session.scalar.return_value = upload_file with ( app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=MagicMock() - ), + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=dataset, + ) as get_dataset_for_tenant, patch( - "controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=MagicMock() - ), + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=MagicMock(), + ) as get_document_by_ref, patch("controllers.console.datasets.datasets_segments.redis_client.setnx", return_value=True), patch( "controllers.console.datasets.datasets_segments.batch_create_segment_to_index_task.delay", @@ -609,6 +612,11 @@ class TestDatasetDocumentSegmentBatchImportApi: ) assert status == 200 assert response["job_status"] == "waiting" + get_dataset_for_tenant.assert_called_once_with("ds-1", "tenant-1", session=session) + document_ref = get_document_by_ref.call_args.args[0] + assert document_ref.dataset.tenant_id == "tenant-1" + assert document_ref.dataset.dataset_id == "ds-1" + assert document_ref.document_id == "doc-1" def test_post_dataset_not_found(self, app: Flask): api = DatasetDocumentSegmentBatchImportApi() @@ -620,7 +628,10 @@ class TestDatasetDocumentSegmentBatchImportApi: with ( app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=None), + patch( + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=None, + ), ): with pytest.raises(NotFound): method( @@ -644,9 +655,13 @@ class TestDatasetDocumentSegmentBatchImportApi: app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=MagicMock() + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=MagicMock(), + ), + patch( + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=None, ), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=None), ): with pytest.raises(NotFound): method( @@ -670,10 +685,12 @@ class TestDatasetDocumentSegmentBatchImportApi: app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=MagicMock() + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=MagicMock(), ), patch( - "controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=MagicMock() + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=MagicMock(), ), ): with pytest.raises(NotFound): @@ -694,10 +711,12 @@ class TestDatasetDocumentSegmentBatchImportApi: app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=MagicMock() + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=MagicMock(), ), patch( - "controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=MagicMock() + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=MagicMock(), ), ): with pytest.raises(ValueError): @@ -718,10 +737,12 @@ class TestDatasetDocumentSegmentBatchImportApi: app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=MagicMock() + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=MagicMock(), ), patch( - "controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=MagicMock() + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=MagicMock(), ), patch( "controllers.console.datasets.datasets_segments.redis_client.setnx", side_effect=Exception("redis down") @@ -1155,8 +1176,14 @@ class TestSegmentOperationCases: with ( app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=None), + patch( + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=dataset, + ), + patch( + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=None, + ), ): with pytest.raises(NotFound): method( @@ -1183,8 +1210,14 @@ class TestSegmentOperationCases: with ( app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), + patch( + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=dataset, + ), + patch( + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=document, + ), ): with pytest.raises(NotFound): method( @@ -1217,11 +1250,13 @@ class TestSegmentOperationCases: with ( app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, + "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", + return_value=dataset, + ), + patch( + "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", + return_value=document, ), patch( "controllers.console.datasets.datasets_segments.batch_create_segment_to_index_task.delay", diff --git a/api/tests/unit_tests/controllers/console/explore/test_audio.py b/api/tests/unit_tests/controllers/console/explore/test_audio.py index 704b45698b6..b2b94440c21 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_audio.py +++ b/api/tests/unit_tests/controllers/console/explore/test_audio.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch import pytest from flask import Flask +from sqlalchemy.orm import Session, scoped_session, sessionmaker from werkzeug.exceptions import InternalServerError import controllers.console.explore.audio as audio_module @@ -22,6 +23,8 @@ from core.errors.error import ( QuotaExceededError, ) from graphon.model_runtime.errors.invoke import InvokeError +from models import Account +from models.model import App, AppMode, InstalledApp from services.app_ref_service import AppRef, MessageRef from services.errors.audio import ( AudioTooLargeServiceError, @@ -40,13 +43,33 @@ def unwrap(func): @pytest.fixture -def installed_app(): - app = MagicMock() - app.app = MagicMock() - app.app.id = "app-1" - app.app.tenant_id = "tenant-1" - app.app_with_session.return_value = app.app - return app +def installed_app( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +): + app = App( + id="app-1", + tenant_id="tenant-1", + name="Explore App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + ) + installed = InstalledApp( + tenant_id="viewer-tenant", + app_id=app.id, + app_owner_tenant_id=app.tenant_id, + position=0, + is_pinned=False, + last_used_at=None, + ) + sqlite_session.add_all([app, installed]) + sqlite_session.commit() + session_proxy = scoped_session(sqlite_session_factory) + monkeypatch.setattr(audio_module.db, "session", session_proxy) + yield installed + session_proxy.remove() @pytest.fixture @@ -259,6 +282,8 @@ class TestChatTextApi: self.method = unwrap(self.api.post) def test_post_success(self, app: Flask, installed_app): + account = Account(name="User", email="user@example.com") + account.id = "account-1" transcript_tts = MagicMock(return_value={"audio": "ok"}) with ( @@ -269,7 +294,7 @@ class TestChatTextApi: patch.object( audio_module, "current_account_with_tenant", - return_value=(MagicMock(id="account-1"), "tenant-1"), + return_value=(account, "tenant-1"), ), patch.object(audio_module.AudioService, "transcript_tts", transcript_tts), ): diff --git a/api/tests/unit_tests/controllers/console/explore/test_completion.py b/api/tests/unit_tests/controllers/console/explore/test_completion.py index 9ff013e27a2..a3a307ae264 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_completion.py +++ b/api/tests/unit_tests/controllers/console/explore/test_completion.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, PropertyMock, patch import pytest from flask import Flask -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, object_session from werkzeug.exceptions import InternalServerError import controllers.console.explore.completion as completion_module @@ -14,7 +14,7 @@ from controllers.console.app.error import ( from controllers.console.explore.error import NotChatAppError, NotCompletionAppError from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError from models import Account -from models.model import AppMode +from models.model import App, AppMode, InstalledApp from services.errors.llm import InvokeRateLimitError @@ -25,23 +25,50 @@ def user(): return account -@pytest.fixture -def completion_app(): - return _installed_app(AppMode.COMPLETION) +@pytest.fixture(autouse=True) +def bind_database(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(completion_module.db, "session", sqlite_session) @pytest.fixture -def chat_app(): - return _installed_app(AppMode.CHAT) +def completion_app(sqlite_session: Session) -> InstalledApp: + return _installed_app(AppMode.COMPLETION, sqlite_session) -def _installed_app(mode: AppMode): - app = MagicMock(mode=mode) - installed_app = MagicMock(app=app) - installed_app.app_with_session.return_value = app +@pytest.fixture +def chat_app(sqlite_session: Session) -> InstalledApp: + return _installed_app(AppMode.CHAT, sqlite_session) + + +def _installed_app(mode: AppMode, session: Session) -> InstalledApp: + app = App( + tenant_id="owner-tenant", + name=f"{mode.value} App", + mode=mode, + enable_site=True, + enable_api=False, + ) + session.add(app) + session.flush() + installed_app = InstalledApp( + tenant_id="viewer-tenant", + app_id=app.id, + app_owner_tenant_id=app.tenant_id, + position=0, + is_pinned=False, + last_used_at=None, + ) + session.add(installed_app) + session.commit() return installed_app +def _session(installed_app: InstalledApp) -> Session: + session = object_session(installed_app) + assert session is not None + return session + + @pytest.fixture def payload_data(): return {"inputs": {}, "query": "hi"} @@ -79,24 +106,24 @@ class TestCompletionApi: result = method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) assert result == ("ok", 200) - def test_post_wrong_app_mode(self, user): + def test_post_wrong_app_mode(self, user, sqlite_session: Session): api = completion_module.CompletionApi() method = unwrap(api.post) - installed_app = _installed_app(AppMode.CHAT) + installed_app = _installed_app(AppMode.CHAT, sqlite_session) with pytest.raises(NotCompletionAppError): method( api, completion_module.CompletionMessageExplorePayload.model_validate({"inputs": {}, "query": "hi"}), - MagicMock(), + sqlite_session, user, installed_app, ) @@ -118,7 +145,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -140,7 +167,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -162,7 +189,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -184,7 +211,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -206,7 +233,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -228,7 +255,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -250,7 +277,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -272,7 +299,7 @@ class TestCompletionApi: method( api, completion_module.CompletionMessageExplorePayload.model_validate(payload_data), - MagicMock(), + _session(completion_app), user, completion_app, ) @@ -284,19 +311,19 @@ class TestCompletionStopApi: method = unwrap(api.post) with patch.object(completion_module.AppTaskService, "stop_task"): - resp, status = method(api, MagicMock(), "u1", completion_app, "task-1") + resp, status = method(api, _session(completion_app), "u1", completion_app, "task-1") assert status == 200 assert resp == {"result": "success"} - def test_stop_wrong_app_mode(self): + def test_stop_wrong_app_mode(self, sqlite_session: Session): api = completion_module.CompletionStopApi() method = unwrap(api.post) - installed_app = _installed_app(AppMode.CHAT) + installed_app = _installed_app(AppMode.CHAT, sqlite_session) with pytest.raises(NotCompletionAppError): - method(api, MagicMock(), "u1", installed_app, "task") + method(api, sqlite_session, "u1", installed_app, "task") class TestChatApi: @@ -319,22 +346,26 @@ class TestChatApi: ), ): result = method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) assert result == ("ok", 200) - def test_post_not_chat_app(self, user): + def test_post_not_chat_app(self, user, sqlite_session: Session): api = completion_module.ChatApi() method = unwrap(api.post) - installed_app = _installed_app(AppMode.COMPLETION) + installed_app = _installed_app(AppMode.COMPLETION, sqlite_session) with pytest.raises(NotChatAppError): method( api, completion_module.ChatMessagePayload.model_validate({"inputs": {}, "query": "hi"}), - MagicMock(), + sqlite_session, user, installed_app, ) @@ -354,7 +385,11 @@ class TestChatApi: ): with pytest.raises(InvokeRateLimitHttpError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) def test_conversation_completed_chat(self, app: Flask, chat_app, user, payload_patch, payload_data): @@ -372,7 +407,11 @@ class TestChatApi: ): with pytest.raises(ConversationCompletedError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) def test_conversation_not_exists_chat(self, app: Flask, chat_app, user, payload_patch, payload_data): @@ -390,12 +429,14 @@ class TestChatApi: ): with pytest.raises(completion_module.NotFound): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) - def test_invalid_conversation_id_fails_fast_as_not_found( - self, app: Flask, chat_app, user, unbound_session: Session - ) -> None: + def test_invalid_conversation_id_fails_fast_as_not_found(self, app: Flask, chat_app, user) -> None: # A nonexistent conversation_id must fail fast as 404, before the streaming # generator is created. Previously the lookup only ran inside the generator, # so an invalid id surfaced as a hang instead of a clean error. @@ -410,7 +451,7 @@ class TestChatApi: get_conversation_mock = MagicMock( side_effect=completion_module.services.errors.conversation.ConversationNotExistsError() ) - session = unbound_session + session = _session(chat_app) api = completion_module.ChatApi() method = unwrap(api.post) @@ -455,7 +496,11 @@ class TestChatApi: ): with pytest.raises(completion_module.AppUnavailableError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) def test_provider_not_initialized_chat(self, app: Flask, chat_app, user, payload_patch, payload_data): @@ -473,7 +518,11 @@ class TestChatApi: ): with pytest.raises(completion_module.ProviderNotInitializeError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) def test_quota_exceeded_chat(self, app: Flask, chat_app, user, payload_patch, payload_data): @@ -491,7 +540,11 @@ class TestChatApi: ): with pytest.raises(completion_module.ProviderQuotaExceededError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) def test_model_not_supported_chat(self, app: Flask, chat_app, user, payload_patch, payload_data): @@ -509,7 +562,11 @@ class TestChatApi: ): with pytest.raises(completion_module.ProviderModelCurrentlyNotSupportError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) def test_invoke_error_chat(self, app: Flask, chat_app, user, payload_patch, payload_data): @@ -527,7 +584,11 @@ class TestChatApi: ): with pytest.raises(completion_module.CompletionRequestError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) def test_internal_error_chat(self, app: Flask, chat_app, user, payload_patch, payload_data): @@ -545,7 +606,11 @@ class TestChatApi: ): with pytest.raises(InternalServerError): method( - api, completion_module.ChatMessagePayload.model_validate(payload_data), MagicMock(), user, chat_app + api, + completion_module.ChatMessagePayload.model_validate(payload_data), + _session(chat_app), + user, + chat_app, ) @@ -554,16 +619,16 @@ class TestChatStopApi: api = completion_module.ChatStopApi() method = unwrap(api.post) with patch.object(completion_module.AppTaskService, "stop_task"): - resp, status = method(api, MagicMock(), "u1", chat_app, "task-1") + resp, status = method(api, _session(chat_app), "u1", chat_app, "task-1") assert status == 200 assert resp == {"result": "success"} - def test_stop_not_chat_app(self): + def test_stop_not_chat_app(self, sqlite_session: Session): api = completion_module.ChatStopApi() method = unwrap(api.post) - installed_app = _installed_app(AppMode.COMPLETION) + installed_app = _installed_app(AppMode.COMPLETION, sqlite_session) with pytest.raises(NotChatAppError): - method(api, MagicMock(), "u1", installed_app, "task") + method(api, sqlite_session, "u1", installed_app, "task") diff --git a/api/tests/unit_tests/controllers/console/explore/test_message.py b/api/tests/unit_tests/controllers/console/explore/test_message.py index cb5a50f2346..0f8e37bc5ff 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_message.py +++ b/api/tests/unit_tests/controllers/console/explore/test_message.py @@ -1,7 +1,9 @@ +from decimal import Decimal from unittest.mock import MagicMock, patch import pytest from flask import Flask +from sqlalchemy.orm import Session, scoped_session, sessionmaker from werkzeug.exceptions import InternalServerError, NotFound import controllers.console.explore.message as module @@ -23,6 +25,9 @@ from core.errors.error import ( QuotaExceededError, ) from graphon.model_runtime.errors.invoke import InvokeError +from models import Account +from models.enums import ConversationFromSource +from models.model import App, AppMode, InstalledApp, Message from services.errors.conversation import ConversationNotExistsError from services.errors.message import ( FirstMessageNotExistsError, @@ -40,43 +45,82 @@ def unwrap(func): return func -def make_message(): - msg = MagicMock() - msg.id = "m1" - msg.conversation_id = "11111111-1111-1111-1111-111111111111" - msg.parent_message_id = None - msg.inputs = {} - msg.query = "hello" - msg.re_sign_file_url_answer = "" - msg.user_feedback = MagicMock(rating=None) - msg.inputs_with_session.return_value = msg.inputs - msg.user_feedback_with_session.return_value = msg.user_feedback - msg.total_price = None - msg.currency = None - msg.status = "normal" - msg.error = None - return msg +def make_message(*, app_id: str): + message = Message( + id="m1", + app_id=app_id, + conversation_id="11111111-1111-1111-1111-111111111111", + query="hello", + message={"role": "user", "content": "hello"}, + answer="", + message_tokens=0, + message_unit_price=Decimal(0), + answer_tokens=0, + answer_unit_price=Decimal(0), + provider_response_latency=0, + currency="USD", + from_source=ConversationFromSource.API, + app_mode=AppMode.CHAT, + ) + message._inputs = {} + message.status = "normal" + return message -def make_installed_app(mode: str | None = None): - app_model = MagicMock(mode=mode) - installed_app = MagicMock() - installed_app.app = app_model - installed_app.app_with_session.return_value = app_model +def make_installed_app(session: Session, mode: str | None = None) -> InstalledApp: + app_model = App( + tenant_id="owner-tenant", + name="Explore App", + mode=mode or AppMode.CHAT, + enable_site=True, + enable_api=False, + ) + session.add(app_model) + session.flush() + installed_app = InstalledApp( + tenant_id="viewer-tenant", + app_id=app_model.id, + app_owner_tenant_id=app_model.tenant_id, + position=0, + is_pinned=False, + last_used_at=None, + ) + session.add(installed_app) + session.commit() return installed_app -class TestMessageListApi: +class _UsesSQLiteSession: + sqlite_session: Session + account: Account + + @pytest.fixture(autouse=True) + def _bind_database( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ): + self.sqlite_session = sqlite_session + self.account = Account(name="User", email="user@example.com") + self.account.id = "account-1" + session_proxy = scoped_session(sqlite_session_factory) + monkeypatch.setattr(module.db, "session", session_proxy) + yield + session_proxy.remove() + + +class TestMessageListApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = module.MessageListApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") pagination = MagicMock( limit=20, has_more=False, - data=[make_message(), make_message()], + data=[make_message(app_id=installed_app.app_id), make_message(app_id=installed_app.app_id)], ) with ( @@ -90,7 +134,7 @@ class TestMessageListApi: return_value=pagination, ), ): - result = method(MagicMock(), installed_app) + result = method(self.account, installed_app) assert result["limit"] == 20 assert result["has_more"] is False @@ -100,16 +144,16 @@ class TestMessageListApi: api = module.MessageListApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with pytest.raises(NotChatAppError): - method(MagicMock(), installed_app) + method(self.account, installed_app) def test_conversation_not_exists(self, app: Flask): api = module.MessageListApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( app.test_request_context( @@ -123,13 +167,13 @@ class TestMessageListApi: ), ): with pytest.raises(NotFound): - method(MagicMock(), installed_app) + method(self.account, installed_app) def test_first_message_not_exists(self, app: Flask): api = module.MessageListApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( app.test_request_context( @@ -143,15 +187,15 @@ class TestMessageListApi: ), ): with pytest.raises(NotFound): - method(MagicMock(), installed_app) + method(self.account, installed_app) -class TestMessageFeedbackApi: +class TestMessageFeedbackApi(_UsesSQLiteSession): def test_post_success(self, app: Flask): api = module.MessageFeedbackApi() method = unwrap(api.post) - installed_app = make_installed_app() + installed_app = make_installed_app(self.sqlite_session) with ( app.test_request_context("/", json={"rating": "like"}), @@ -161,7 +205,7 @@ class TestMessageFeedbackApi: ), ): result = method( - module.MessageFeedbackPayload.model_validate({"rating": "like"}), MagicMock(), installed_app, "mid" + module.MessageFeedbackPayload.model_validate({"rating": "like"}), self.account, installed_app, "mid" ) assert result["result"] == "success" @@ -170,7 +214,7 @@ class TestMessageFeedbackApi: api = module.MessageFeedbackApi() method = unwrap(api.post) - installed_app = make_installed_app() + installed_app = make_installed_app(self.sqlite_session) with ( app.test_request_context("/", json={}), @@ -181,15 +225,15 @@ class TestMessageFeedbackApi: ), ): with pytest.raises(NotFound): - method(module.MessageFeedbackPayload.model_validate({}), MagicMock(), installed_app, "mid") + method(module.MessageFeedbackPayload.model_validate({}), self.account, installed_app, "mid") -class TestMessageMoreLikeThisApi: +class TestMessageMoreLikeThisApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -207,7 +251,7 @@ class TestMessageMoreLikeThisApi: return_value=("ok", 200), ), ): - resp = method(MagicMock(), MagicMock(), installed_app, "mid") + resp = method(self.sqlite_session, self.account, installed_app, "mid") assert resp == ("ok", 200) @@ -215,16 +259,16 @@ class TestMessageMoreLikeThisApi: api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with pytest.raises(NotCompletionAppError): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") def test_more_like_this_disabled(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -238,13 +282,13 @@ class TestMessageMoreLikeThisApi: ), ): with pytest.raises(AppMoreLikeThisDisabledError): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") def test_message_not_exists_more_like_this(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -258,13 +302,13 @@ class TestMessageMoreLikeThisApi: ), ): with pytest.raises(NotFound): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") def test_provider_not_init_more_like_this(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -278,13 +322,13 @@ class TestMessageMoreLikeThisApi: ), ): with pytest.raises(ProviderNotInitializeError): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") def test_quota_exceeded_more_like_this(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -298,13 +342,13 @@ class TestMessageMoreLikeThisApi: ), ): with pytest.raises(ProviderQuotaExceededError): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") def test_model_not_support_more_like_this(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -318,13 +362,13 @@ class TestMessageMoreLikeThisApi: ), ): with pytest.raises(ProviderModelCurrentlyNotSupportError): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") def test_invoke_error_more_like_this(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -338,13 +382,13 @@ class TestMessageMoreLikeThisApi: ), ): with pytest.raises(CompletionRequestError): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") def test_unexpected_error_more_like_this(self, app: Flask): api = module.MessageMoreLikeThisApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with ( app.test_request_context( @@ -358,15 +402,15 @@ class TestMessageMoreLikeThisApi: ), ): with pytest.raises(InternalServerError): - method(MagicMock(), MagicMock(), installed_app, "mid") + method(self.sqlite_session, self.account, installed_app, "mid") -class TestMessageSuggestedQuestionApi: +class TestMessageSuggestedQuestionApi(_UsesSQLiteSession): def test_get_success(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -375,7 +419,7 @@ class TestMessageSuggestedQuestionApi: return_value=["q1", "q2"], ), ): - result = method(MagicMock(), installed_app, "mid") + result = method(self.account, installed_app, "mid") assert result["data"] == ["q1", "q2"] @@ -383,16 +427,16 @@ class TestMessageSuggestedQuestionApi: api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode="completion") with pytest.raises(NotChatAppError): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_disabled(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -402,13 +446,13 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(AppSuggestedQuestionsAfterAnswerDisabledError): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_message_not_exists_suggested_question(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -418,13 +462,13 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(NotFound): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_conversation_not_exists_suggested_question(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -434,13 +478,13 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(NotFound): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_provider_not_init_suggested_question(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -450,13 +494,13 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(ProviderNotInitializeError): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_quota_exceeded_suggested_question(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -466,13 +510,13 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(ProviderQuotaExceededError): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_model_not_support_suggested_question(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -482,13 +526,13 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(ProviderModelCurrentlyNotSupportError): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_invoke_error_suggested_question(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -498,13 +542,13 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(CompletionRequestError): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") def test_unexpected_error_suggested_question(self): api = module.MessageSuggestedQuestionApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode="chat") with ( patch.object( @@ -514,4 +558,4 @@ class TestMessageSuggestedQuestionApi: ), ): with pytest.raises(InternalServerError): - method(MagicMock(), installed_app, "mid") + method(self.account, installed_app, "mid") diff --git a/api/tests/unit_tests/controllers/console/explore/test_parameter.py b/api/tests/unit_tests/controllers/console/explore/test_parameter.py index e1a0cd11ac2..a41a3805cc5 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_parameter.py +++ b/api/tests/unit_tests/controllers/console/explore/test_parameter.py @@ -6,11 +6,19 @@ import pytest import controllers.console.explore.parameter as module from controllers.console.app.error import AppUnavailableError +from models.model import InstalledApp from services.app_definition_query_service import AppDefinitionQueryService, AppDefinitionUnavailableError -def _installed_app() -> MagicMock: - return MagicMock(app_id="app-1") +def _installed_app() -> InstalledApp: + return InstalledApp( + tenant_id="viewer-tenant", + app_id="app-1", + app_owner_tenant_id="owner-tenant", + position=0, + is_pinned=False, + last_used_at=None, + ) def _application_services() -> tuple[SimpleNamespace, MagicMock]: diff --git a/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py b/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py index d354e22541d..fccbba38d6d 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py +++ b/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py @@ -1,14 +1,24 @@ from inspect import unwrap -from unittest.mock import ANY, patch +from types import SimpleNamespace +from unittest.mock import MagicMock, patch import pytest from flask import Flask from pydantic import ValidationError import controllers.console.explore.recommended_app as module -from controllers.console.explore.recommended_app import RecommendedAppsQuery from models import Account from models.model import AppMode, IconType +from services.recommended_app_query_service import ( + LearnDifyAppListResult, + RecommendedAppDetailSummary, + RecommendedAppInfoRecord, + RecommendedAppListResult, + RecommendedAppSummary, +) +from services.recommended_app_query_service import ( + RecommendedAppNotFoundError as RecommendedAppQueryNotFoundError, +) def make_account(interface_language: str | None) -> Account: @@ -19,133 +29,149 @@ def make_account(interface_language: str | None) -> Account: class TestRecommendedAppListApi: - def test_get_with_language_param(self, app: Flask): + def test_get_with_language_param(self, app: Flask) -> None: api = module.RecommendedAppListApi() method = unwrap(api.get) - result_data = {"recommended_apps": [], "categories": []} + queries = MagicMock() + queries.list_recommended.return_value = RecommendedAppListResult(recommended_apps=(), categories=()) with ( app.test_request_context("/", query_string={"language": "en-US"}), patch.object( - module.RecommendedAppService, - "get_recommended_apps_and_categories", - return_value=result_data, - ) as service_mock, + module, + "application_services", + return_value=SimpleNamespace(recommended_app_queries=queries), + ), ): - result = method(api, RecommendedAppsQuery(language="en-US"), make_account("fr-FR")) + result = method(api, module.RecommendedAppsQuery(language="en-US"), make_account("fr-FR")) - service_mock.assert_called_once_with("en-US", session=ANY) - assert result == result_data + queries.list_recommended.assert_called_once_with( + requested_language="en-US", + interface_language="fr-FR", + ) + assert result == {"recommended_apps": [], "categories": []} - def test_get_fallback_to_user_language(self, app: Flask): - api = module.RecommendedAppListApi() + +class TestLearnDifyAppListApi: + def test_get_with_language_param(self, app: Flask) -> None: + api = module.LearnDifyAppListApi() method = unwrap(api.get) - result_data = {"recommended_apps": [], "categories": []} + queries = MagicMock() + queries.list_learn_dify.return_value = LearnDifyAppListResult(recommended_apps=()) with ( - app.test_request_context("/", query_string={"language": "invalid"}), + app.test_request_context("/", query_string={"language": "en-US"}), patch.object( - module.RecommendedAppService, - "get_recommended_apps_and_categories", - return_value=result_data, - ) as service_mock, + module, + "application_services", + return_value=SimpleNamespace(recommended_app_queries=queries), + ), ): - result = method(api, RecommendedAppsQuery(), make_account("fr-FR")) + result = method(api, module.RecommendedAppsQuery(language="en-US"), make_account("fr-FR")) - service_mock.assert_called_once_with("fr-FR", session=ANY) - assert result == result_data + queries.list_learn_dify.assert_called_once_with( + requested_language="en-US", + interface_language="fr-FR", + ) + assert result == {"recommended_apps": []} - def test_get_fallback_to_default_language(self, app: Flask): - api = module.RecommendedAppListApi() + +class TestRecommendedAppApi: + def test_get_success(self, app: Flask) -> None: + api = module.RecommendedAppApi() method = unwrap(api.get) - result_data = {"recommended_apps": [], "categories": []} + queries = MagicMock() + queries.get_detail.return_value = RecommendedAppDetailSummary( + id="app1", + name="App", + icon=None, + icon_background=None, + mode="chat", + export_data="{}", + can_trial=False, + ) with ( app.test_request_context("/"), patch.object( - module.RecommendedAppService, - "get_recommended_apps_and_categories", - return_value=result_data, - ) as service_mock, + module, + "application_services", + return_value=SimpleNamespace(recommended_app_queries=queries), + ), ): - result = method(api, RecommendedAppsQuery(), make_account(None)) + result = method(api, "11111111-1111-1111-1111-111111111111") - service_mock.assert_called_once_with(module.languages[0], session=ANY) - assert result == result_data - - -class TestLearnDifyAppListApi: - def test_get_with_language_param(self, app: Flask): - api = module.LearnDifyAppListApi() - method = unwrap(api.get) - - result_data = {"recommended_apps": []} - - with ( - app.test_request_context("/", query_string={"language": "en-US"}), - patch.object( - module.RecommendedAppService, - "get_learn_dify_apps", - return_value=result_data, - ) as service_mock, - ): - result = method(api, RecommendedAppsQuery(language="en-US"), make_account("fr-FR")) - - service_mock.assert_called_once_with("en-US", session=ANY) - assert result == result_data - - def test_get_fallback_to_user_language(self, app: Flask): - api = module.LearnDifyAppListApi() - method = unwrap(api.get) - - result_data = {"recommended_apps": []} - - with ( - app.test_request_context("/", query_string={"language": "invalid"}), - patch.object( - module.RecommendedAppService, - "get_learn_dify_apps", - return_value=result_data, - ) as service_mock, - ): - result = method(api, RecommendedAppsQuery(), make_account("fr-FR")) - - service_mock.assert_called_once_with("fr-FR", session=ANY) - assert result == result_data - - -class TestRecommendedAppApi: - def test_get_success(self, app: Flask): - api = module.RecommendedAppApi() - method = unwrap(api.get) - - result_data = { + queries.get_detail.assert_called_once_with("11111111-1111-1111-1111-111111111111") + assert result == { "id": "app1", "name": "App", + "icon": None, + "icon_background": None, "mode": "chat", "export_data": "{}", "can_trial": False, } + def test_get_missing_raises_stable_not_found_error(self, app: Flask) -> None: + api = module.RecommendedAppApi() + method = unwrap(api.get) + queries = MagicMock() + queries.get_detail.side_effect = RecommendedAppQueryNotFoundError + with ( app.test_request_context("/"), patch.object( - module.RecommendedAppService, - "get_recommend_app_detail", - return_value=result_data, - ) as service_mock, + module, + "application_services", + return_value=SimpleNamespace(recommended_app_queries=queries), + ), ): - result = method(api, "11111111-1111-1111-1111-111111111111") + with pytest.raises(module.RecommendedAppNotFoundError) as exc_info: + method(api, "11111111-1111-1111-1111-111111111111") - service_mock.assert_called_once_with("11111111-1111-1111-1111-111111111111", session=ANY) - assert result == {**result_data, "icon": None, "icon_background": None} + assert exc_info.value.data == { + "code": "recommended_app_not_found", + "message": "Recommended app not found.", + "status": 404, + } class TestRecommendedAppResponseModels: - def test_recommended_app_info_response_computes_icon_url(self): + def test_query_service_records_serialize_through_controller_contract(self) -> None: + result = RecommendedAppListResult( + recommended_apps=( + RecommendedAppSummary( + app=RecommendedAppInfoRecord( + id="app-1", + name="App", + mode="chat", + icon=None, + icon_type=None, + icon_background=None, + ), + app_id="app-1", + description=None, + copyright=None, + privacy_policy=None, + custom_disclaimer=None, + categories=("Workflow",), + position=1, + is_listed=True, + can_trial=False, + ), + ), + categories=("Workflow",), + ) + + response = module.dump_response(module.RecommendedAppListResponse, result) + + assert response["recommended_apps"][0]["app"]["id"] == "app-1" + assert response["recommended_apps"][0]["categories"] == ["Workflow"] + + def test_recommended_app_info_response_computes_icon_url(self) -> None: with patch.object(module, "build_icon_url", return_value="https://signed/icon.png"): payload = module.RecommendedAppInfoResponse.model_validate( { @@ -160,7 +186,7 @@ class TestRecommendedAppResponseModels: assert payload["icon_url"] == "https://signed/icon.png" - def test_recommended_app_list_response_serialization(self): + def test_recommended_app_list_response_serialization(self) -> None: response = module.RecommendedAppListResponse.model_validate( { "recommended_apps": [ @@ -189,7 +215,7 @@ class TestRecommendedAppResponseModels: assert response["recommended_apps"][0]["categories"] == ["cat", "other"] assert response["categories"] == ["cat"] - def test_learn_dify_app_list_response_serialization(self): + def test_learn_dify_app_list_response_serialization(self) -> None: response = module.LearnDifyAppListResponse.model_validate( { "recommended_apps": [ @@ -216,6 +242,6 @@ class TestRecommendedAppResponseModels: assert response["recommended_apps"][0]["app_id"] == "app-1" assert response["recommended_apps"][0]["categories"] == ["Workflow"] - def test_recommended_app_response_requires_can_trial(self): + def test_recommended_app_response_requires_can_trial(self) -> None: with pytest.raises(ValidationError): module.RecommendedAppResponse.model_validate({"app_id": "app-1"}) diff --git a/api/tests/unit_tests/controllers/console/explore/test_saved_message.py b/api/tests/unit_tests/controllers/console/explore/test_saved_message.py index a685c6c8fbf..3ca6b9d179b 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_saved_message.py +++ b/api/tests/unit_tests/controllers/console/explore/test_saved_message.py @@ -1,36 +1,77 @@ +from decimal import Decimal from inspect import unwrap from unittest.mock import MagicMock, PropertyMock, patch from uuid import uuid4 import pytest from flask import Flask +from sqlalchemy.orm import Session, scoped_session, sessionmaker from werkzeug.exceptions import NotFound import controllers.console.explore.saved_message as module from controllers.console.explore.error import NotCompletionAppError +from models import Account +from models.enums import ConversationFromSource, FeedbackFromSource, FeedbackRating +from models.model import App, AppMode, InstalledApp, Message, MessageFeedback from services.errors.message import MessageNotExistsError -def make_saved_message(): - msg = MagicMock() - msg.id = str(uuid4()) - msg.message_id = str(uuid4()) - msg.app_id = str(uuid4()) - msg.inputs = {} - msg.query = "hello" - msg.answer = "world" - msg.user_feedback = MagicMock(rating="like") - msg.inputs_with_session.return_value = msg.inputs - msg.user_feedback_with_session.return_value = msg.user_feedback - msg.created_at = None - return msg +def make_message(session: Session, *, app_id: str, account_id: str) -> Message: + message = Message( + id=str(uuid4()), + app_id=app_id, + conversation_id=str(uuid4()), + query="hello", + message={"role": "user", "content": "hello"}, + answer="world", + message_tokens=1, + message_unit_price=Decimal(0), + answer_tokens=1, + answer_unit_price=Decimal(0), + provider_response_latency=0, + currency="USD", + from_source=ConversationFromSource.API, + from_account_id=account_id, + app_mode=AppMode.COMPLETION, + ) + message._inputs = {} + message.status = "normal" + session.add(message) + session.flush() + session.add( + MessageFeedback( + app_id=app_id, + conversation_id=message.conversation_id, + message_id=message.id, + rating=FeedbackRating.LIKE, + from_source=FeedbackFromSource.USER, + from_account_id=account_id, + ) + ) + session.flush() + return message -def make_installed_app(mode: str): - app_model = MagicMock(mode=mode) - installed_app = MagicMock() - installed_app.app = app_model - installed_app.app_with_session.return_value = app_model +def make_installed_app(session: Session, mode: AppMode | str) -> InstalledApp: + app_model = App( + tenant_id="owner-tenant", + name="Explore App", + mode=AppMode.value_of(mode), + enable_site=True, + enable_api=False, + ) + session.add(app_model) + session.flush() + installed_app = InstalledApp( + tenant_id="viewer-tenant", + app_id=app_model.id, + app_owner_tenant_id=app_model.tenant_id, + position=0, + is_pinned=False, + last_used_at=None, + ) + session.add(installed_app) + session.commit() return installed_app @@ -47,19 +88,43 @@ def payload_patch(): return _patch -class TestSavedMessageListApi: +class _UsesSQLiteSession: + sqlite_session: Session + account: Account + + @pytest.fixture(autouse=True) + def _bind_database( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ): + self.sqlite_session = sqlite_session + self.account = Account(name="User", email="user@example.com") + self.sqlite_session.add(self.account) + self.sqlite_session.commit() + session_proxy = scoped_session(sqlite_session_factory) + monkeypatch.setattr(module.db, "session", session_proxy) + yield + session_proxy.remove() + + +class TestSavedMessageListApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = module.SavedMessageListApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode=AppMode.COMPLETION) pagination = MagicMock( limit=20, has_more=False, - data=[make_saved_message(), make_saved_message()], + data=[ + make_message(self.sqlite_session, app_id=installed_app.app_id, account_id=self.account.id), + make_message(self.sqlite_session, app_id=installed_app.app_id, account_id=self.account.id), + ], ) - current_user = MagicMock() + self.sqlite_session.commit() with ( app.test_request_context("/", query_string={}), @@ -69,10 +134,10 @@ class TestSavedMessageListApi: return_value=pagination, ) as pagination_mock, ): - result = method(api, current_user, installed_app) + result = method(api, self.account, installed_app) pagination_mock.assert_called_once() - assert pagination_mock.call_args.args[1] is current_user + assert pagination_mock.call_args.args[1] is self.account assert result["limit"] == 20 assert result["has_more"] is False assert len(result["data"]) == 2 @@ -81,36 +146,35 @@ class TestSavedMessageListApi: api = module.SavedMessageListApi() method = unwrap(api.get) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode=AppMode.CHAT) with pytest.raises(NotCompletionAppError): - method(api, MagicMock(), installed_app) + method(api, self.account, installed_app) def test_post_success(self, app: Flask, payload_patch): api = module.SavedMessageListApi() method = unwrap(api.post) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode=AppMode.COMPLETION) payload = {"message_id": str(uuid4())} - current_user = MagicMock() with ( app.test_request_context("/", json=payload), payload_patch(payload), patch.object(module.SavedMessageService, "save") as save_mock, ): - result = method(api, module.SavedMessageCreatePayload.model_validate(payload), current_user, installed_app) + result = method(api, module.SavedMessageCreatePayload.model_validate(payload), self.account, installed_app) save_mock.assert_called_once() - assert save_mock.call_args.args[1] is current_user + assert save_mock.call_args.args[1] is self.account assert result == {"result": "success"} def test_post_message_not_exists(self, app: Flask, payload_patch): api = module.SavedMessageListApi() method = unwrap(api.post) - installed_app = make_installed_app(mode="completion") + installed_app = make_installed_app(self.sqlite_session, mode=AppMode.COMPLETION) payload = {"message_id": str(uuid4())} @@ -124,24 +188,23 @@ class TestSavedMessageListApi: ), ): with pytest.raises(NotFound): - method(api, module.SavedMessageCreatePayload.model_validate(payload), MagicMock(), installed_app) + method(api, module.SavedMessageCreatePayload.model_validate(payload), self.account, installed_app) -class TestSavedMessageApi: +class TestSavedMessageApi(_UsesSQLiteSession): def test_delete_success(self): api = module.SavedMessageApi() method = unwrap(api.delete) - installed_app = make_installed_app(mode="completion") - current_user = MagicMock() + installed_app = make_installed_app(self.sqlite_session, mode=AppMode.COMPLETION) with ( patch.object(module.SavedMessageService, "delete") as delete_mock, ): - result, status = method(api, current_user, installed_app, str(uuid4())) + result, status = method(api, self.account, installed_app, str(uuid4())) delete_mock.assert_called_once() - assert delete_mock.call_args.args[1] is current_user + assert delete_mock.call_args.args[1] is self.account assert status == 204 assert result == "" @@ -149,7 +212,7 @@ class TestSavedMessageApi: api = module.SavedMessageApi() method = unwrap(api.delete) - installed_app = make_installed_app(mode="chat") + installed_app = make_installed_app(self.sqlite_session, mode=AppMode.CHAT) with pytest.raises(NotCompletionAppError): - method(api, MagicMock(), installed_app, str(uuid4())) + method(api, self.account, installed_app, str(uuid4())) diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial.py b/api/tests/unit_tests/controllers/console/explore/test_trial.py index 2024ec9473d..2f2966b3844 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_trial.py +++ b/api/tests/unit_tests/controllers/console/explore/test_trial.py @@ -9,7 +9,6 @@ from uuid import uuid4 import pytest from flask import Flask, request -from sqlalchemy.engine import Engine from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, InternalServerError, NotFound @@ -40,9 +39,12 @@ from core.helper import encrypter from core.workflow.llm_environment_variable import LLMEnvironmentVariable from graphon.model_runtime.errors.invoke import InvokeError from graphon.variables import SecretVariable, StringVariable -from models import Account +from models import Account, Tenant from models.account import TenantStatus -from models.model import AppMode, Site +from models.dataset import Dataset +from models.model import App, AppMode, Site, UploadFile +from models.tools import WorkflowToolProvider +from models.workflow import Workflow from services.app_ref_service import AppRef, MessageRef from services.errors.audio import SpeechToTextDisabledServiceError from services.errors.conversation import ConversationNotExistsError @@ -55,10 +57,8 @@ class _UsesSQLiteSession: sqlite_session: Session @pytest.fixture(autouse=True) - def _provide_sqlite_session(self, sqlite_engine: Engine): - with Session(sqlite_engine, expire_on_commit=False) as session: - self.sqlite_session = session - yield + def _provide_sqlite_session(self, sqlite_session: Session): + self.sqlite_session = sqlite_session @pytest.fixture @@ -68,6 +68,46 @@ def account() -> Account: return acc +@pytest.fixture +def trial_app_usage(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + usage = MagicMock() + monkeypatch.setattr( + module, + "application_services", + MagicMock(return_value=SimpleNamespace(trial_app_usage=usage)), + ) + return usage + + +def _app(*, app_id: str, mode: AppMode, tenant_id: str = "tenant-1") -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name="Trial App", + mode=mode, + enable_site=True, + enable_api=False, + ) + + +def _upload_file(*, file_id: str = "upload-file-id", tenant_id: str = "app-tenant-id") -> UploadFile: + upload_file = UploadFile( + tenant_id=tenant_id, + storage_type="opendal", + key="trial/file.txt", + name="file.txt", + size=1, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="u1", + created_at=datetime(2024, 1, 1), + used=False, + ) + upload_file.id = file_id + return upload_file + + def _file_data() -> Any: file_data: Any = BytesIO(b"fake audio data") file_data.filename = "test.wav" @@ -87,27 +127,18 @@ def _persist_site(sqlite_session: Session, app_id: str) -> Site: @pytest.fixture -def trial_app_chat() -> MagicMock: - app = MagicMock() - app.id = "a-chat" - app.mode = AppMode.CHAT - return app +def trial_app_chat() -> App: + return _app(app_id="a-chat", mode=AppMode.CHAT) @pytest.fixture -def trial_app_completion() -> MagicMock: - app = MagicMock() - app.id = "a-comp" - app.mode = AppMode.COMPLETION - return app +def trial_app_completion() -> App: + return _app(app_id="a-comp", mode=AppMode.COMPLETION) @pytest.fixture -def trial_app_workflow() -> MagicMock: - app = MagicMock() - app.id = "a-workflow" - app.mode = AppMode.WORKFLOW - return app +def trial_app_workflow() -> App: + return _app(app_id="a-workflow", mode=AppMode.WORKFLOW) def test_trial_workflow_uses_trial_scoped_simple_account_model() -> None: @@ -116,30 +147,27 @@ def test_trial_workflow_uses_trial_scoped_simple_account_model() -> None: def test_trial_dataset_list_preserves_slim_dataset_fields(app: Flask, unbound_session: Session): - class DatasetListItem: - id = "dataset-1" - name = "Dataset" - description = "description" - permission = "only_me" - data_source_type = "upload_file" - indexing_technique = "high_quality" - created_by = "user-1" - created_at = datetime(2024, 1, 1, tzinfo=UTC) - permission_keys = ["dataset.acl.readonly"] - - @property - def app_count(self): - raise AssertionError("trial dataset list should not serialize detail-only computed fields") - api = module.DatasetListApi() method = unwrap(api.get) - app_model = SimpleNamespace(tenant_id="tenant-1") + app_model = _app(app_id="app-1", mode=AppMode.CHAT) + dataset = Dataset( + id="dataset-1", + tenant_id=app_model.tenant_id, + name="Dataset", + description="description", + permission="only_me", + data_source_type="upload_file", + indexing_technique="high_quality", + created_by="user-1", + created_at=datetime(2024, 1, 1, tzinfo=UTC), + ) + dataset.permission_keys = ["dataset.acl.readonly"] # type: ignore[attr-defined] with ( app.test_request_context("/?page=1&limit=20&ids=dataset-1"), patch.object( module.DatasetService, "get_datasets_by_ids", - return_value=([DatasetListItem()], 1), + return_value=([dataset], 1), ) as get_datasets, ): result = method(api, unbound_session, app_model) @@ -170,17 +198,17 @@ def test_trial_dataset_list_preserves_slim_dataset_fields(app: Flask, unbound_se "api_type", [module.TrialSitApi, module.TrialAppParameterApi, module.AppApi, module.AppWorkflowApi, module.DatasetListApi], ) -def test_trial_app_handlers_use_explicit_read_session(api_type: type) -> None: +def test_preview_handlers_use_explicit_read_session(api_type: type) -> None: source = getsource(api_type.get) - assert "@with_session(write=False)\n @get_app_model_with_trial(None)" in source + assert "@with_session(write=False)\n @get_previewable_app_model(None)" in source assert tuple(signature(api_type.get).parameters)[:3] == ("self", "session", "app_model") def test_trial_app_detail_serializes_with_explicit_session( app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session ) -> None: - app_model = MagicMock() + app_model = _app(app_id="app-1", mode=AppMode.CHAT) response_view = MagicMock() get_app = MagicMock(return_value=app_model) build_view = MagicMock(return_value=response_view) @@ -203,8 +231,8 @@ class TestTrialAppFileUploadApi: def test_upload_uses_trial_app_tenant(self, app: Flask, account: Account) -> None: api = module.TrialAppFileUploadApi() method = unwrap(api.post) - app_model = SimpleNamespace(tenant_id="app-tenant-id") - upload_file = MagicMock() + app_model = _app(app_id="app-1", mode=AppMode.CHAT, tenant_id="app-tenant-id") + upload_file = _upload_file() with ( app.test_request_context("/", method="POST"), @@ -222,7 +250,7 @@ class TestTrialAppRemoteFileUploadApi: def test_upload_uses_trial_app_tenant(self, app: Flask, account: Account) -> None: api = module.TrialAppRemoteFileUploadApi() method = unwrap(api.post) - app_model = SimpleNamespace(tenant_id="app-tenant-id") + app_model = _app(app_id="app-1", mode=AppMode.CHAT, tenant_id="app-tenant-id") remote_file = MagicMock() remote_file.model_dump.return_value = {"id": "upload-file-id"} @@ -249,17 +277,22 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): WorkflowRunRequest.model_validate(request.get_json()), self.sqlite_session, account, - MagicMock(mode=AppMode.CHAT), + _app(app_id="not-workflow", mode=AppMode.CHAT), ) - def test_success(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_success( + self, + app: Flask, + trial_app_workflow: App, + account: Account, + trial_app_usage: MagicMock, + ) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) with ( app.test_request_context("/", json={"inputs": {}}), patch.object(module.AppGenerateService, "generate", return_value=MagicMock()), - patch.object(module.RecommendedAppService, "add_trial_app_record"), ): result = method( api, @@ -270,8 +303,9 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): ) assert result is not None + trial_app_usage.record.assert_called_once_with(app_id="a-workflow", account_id="u1") - def test_workflow_provider_not_init(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_workflow_provider_not_init(self, app: Flask, trial_app_workflow: App, account: Account) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) @@ -292,7 +326,7 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): trial_app_workflow, ) - def test_workflow_quota_exceeded(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_workflow_quota_exceeded(self, app: Flask, trial_app_workflow: App, account: Account) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) @@ -313,7 +347,7 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): trial_app_workflow, ) - def test_workflow_model_not_support(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_workflow_model_not_support(self, app: Flask, trial_app_workflow: App, account: Account) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) @@ -334,7 +368,7 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): trial_app_workflow, ) - def test_workflow_invoke_error(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_workflow_invoke_error(self, app: Flask, trial_app_workflow: App, account: Account) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) @@ -355,7 +389,7 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): trial_app_workflow, ) - def test_workflow_rate_limit_error(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_workflow_rate_limit_error(self, app: Flask, trial_app_workflow: App, account: Account) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) @@ -376,7 +410,7 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): trial_app_workflow, ) - def test_workflow_value_error(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_workflow_value_error(self, app: Flask, trial_app_workflow: App, account: Account) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) @@ -397,7 +431,7 @@ class TestTrialAppWorkflowRunApi(_UsesSQLiteSession): trial_app_workflow, ) - def test_workflow_generic_exception(self, app: Flask, trial_app_workflow: MagicMock, account: Account) -> None: + def test_workflow_generic_exception(self, app: Flask, trial_app_workflow: App, account: Account) -> None: api = module.TrialAppWorkflowRunApi() method = unwrap(api.post) @@ -431,25 +465,31 @@ class TestTrialChatApi(_UsesSQLiteSession): ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, - MagicMock(mode="completion"), + _app(app_id="not-chat", mode=AppMode.COMPLETION), ) - def test_success(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_success( + self, + app: Flask, + trial_app_chat: App, + account: Account, + trial_app_usage: MagicMock, + ) -> None: api = module.TrialChatApi() method = unwrap(api.post) with ( app.test_request_context("/", json={"inputs": {}, "query": "hi"}), patch.object(module.AppGenerateService, "generate", return_value=MagicMock()), - patch.object(module.RecommendedAppService, "add_trial_app_record"), ): result = method( api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) assert result is not None + trial_app_usage.record.assert_called_once_with(app_id="a-chat", account_id="u1") - def test_chat_conversation_not_exists(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_conversation_not_exists(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -466,7 +506,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_conversation_completed(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_conversation_completed(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -483,7 +523,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_app_config_broken(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_app_config_broken(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -500,7 +540,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_provider_not_init(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_provider_not_init(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -517,7 +557,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_quota_exceeded(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_quota_exceeded(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -534,7 +574,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_model_not_support(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_model_not_support(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -551,7 +591,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_invoke_error(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_invoke_error(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -568,7 +608,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_rate_limit_error(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_rate_limit_error(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -585,7 +625,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_value_error(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_value_error(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -602,7 +642,7 @@ class TestTrialChatApi(_UsesSQLiteSession): api, ChatRequest.model_validate(request.get_json()), self.sqlite_session, account, trial_app_chat ) - def test_chat_generic_exception(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_chat_generic_exception(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatApi() method = unwrap(api.post) @@ -632,17 +672,22 @@ class TestTrialCompletionApi(_UsesSQLiteSession): CompletionRequest.model_validate(request.get_json()), self.sqlite_session, account, - MagicMock(mode=AppMode.CHAT), + _app(app_id="not-completion", mode=AppMode.CHAT), ) - def test_success(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_success( + self, + app: Flask, + trial_app_completion: App, + account: Account, + trial_app_usage: MagicMock, + ) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) with ( app.test_request_context("/", json={"inputs": {}, "query": ""}), patch.object(module.AppGenerateService, "generate", return_value=MagicMock()), - patch.object(module.RecommendedAppService, "add_trial_app_record"), ): result = method( api, @@ -653,8 +698,9 @@ class TestTrialCompletionApi(_UsesSQLiteSession): ) assert result is not None + trial_app_usage.record.assert_called_once_with(app_id="a-comp", account_id="u1") - def test_completion_app_config_broken(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_app_config_broken(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -675,7 +721,7 @@ class TestTrialCompletionApi(_UsesSQLiteSession): trial_app_completion, ) - def test_completion_provider_not_init(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_provider_not_init(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -696,7 +742,7 @@ class TestTrialCompletionApi(_UsesSQLiteSession): trial_app_completion, ) - def test_completion_quota_exceeded(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_quota_exceeded(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -717,7 +763,7 @@ class TestTrialCompletionApi(_UsesSQLiteSession): trial_app_completion, ) - def test_completion_model_not_support(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_model_not_support(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -738,7 +784,7 @@ class TestTrialCompletionApi(_UsesSQLiteSession): trial_app_completion, ) - def test_completion_invoke_error(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_invoke_error(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -759,7 +805,7 @@ class TestTrialCompletionApi(_UsesSQLiteSession): trial_app_completion, ) - def test_completion_rate_limit_error(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_rate_limit_error(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -780,7 +826,7 @@ class TestTrialCompletionApi(_UsesSQLiteSession): trial_app_completion, ) - def test_completion_value_error(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_value_error(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -801,7 +847,7 @@ class TestTrialCompletionApi(_UsesSQLiteSession): trial_app_completion, ) - def test_completion_generic_exception(self, app: Flask, trial_app_completion: MagicMock, account: Account) -> None: + def test_completion_generic_exception(self, app: Flask, trial_app_completion: App, account: Account) -> None: api = module.TrialCompletionApi() method = unwrap(api.post) @@ -830,9 +876,9 @@ class TestTrialMessageSuggestedQuestionApi: with app.test_request_context("/"): with pytest.raises(NotChatAppError): - method(api, account, MagicMock(mode="completion"), str(uuid4())) + method(api, account, _app(app_id="not-chat", mode=AppMode.COMPLETION), str(uuid4())) - def test_success(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_success(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialMessageSuggestedQuestionApi() method = unwrap(api.get) @@ -848,7 +894,7 @@ class TestTrialMessageSuggestedQuestionApi: assert result == {"data": ["q1", "q2"]} - def test_conversation_not_exists(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_conversation_not_exists(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialMessageSuggestedQuestionApi() method = unwrap(api.get) @@ -882,7 +928,7 @@ class TestTrialAppParameterApi: services = SimpleNamespace(app_definitions=app_definitions) with patch.object(module, "application_services", return_value=services): - result = method(api, unbound_session, SimpleNamespace(id="app-1")) + result = method(api, unbound_session, _app(app_id="app-1", mode=AppMode.CHAT)) assert result == expected app_definitions.get_parameters.assert_called_once_with("app-1") @@ -898,11 +944,17 @@ class TestTrialAppParameterApi: patch.object(module, "application_services", return_value=services), pytest.raises(AppUnavailableError), ): - method(api, unbound_session, SimpleNamespace(id="app-1")) + method(api, unbound_session, _app(app_id="app-1", mode=AppMode.CHAT)) class TestTrialChatAudioApi: - def test_success(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_success( + self, + app: Flask, + trial_app_chat: App, + account: Account, + trial_app_usage: MagicMock, + ) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -913,13 +965,13 @@ class TestTrialChatAudioApi: "/", method="POST", data={"file": (file_data, "test.wav")}, content_type="multipart/form-data" ), patch.object(module.AudioService, "transcript_asr", return_value={"text": "hello"}), - patch.object(module.RecommendedAppService, "add_trial_app_record"), ): result = method(api, account, trial_app_chat) assert result == {"text": "hello"} + trial_app_usage.record.assert_called_once_with(app_id="a-chat", account_id="u1") - def test_app_config_broken(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_app_config_broken(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -942,7 +994,7 @@ class TestTrialChatAudioApi: trial_app_chat, ) - def test_no_audio_uploaded(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_no_audio_uploaded(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -965,7 +1017,7 @@ class TestTrialChatAudioApi: trial_app_chat, ) - def test_missing_file_field_returns_400(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_missing_file_field_returns_400(self, app: Flask, trial_app_chat: App, account: Account) -> None: """A multipart POST with no `file` field must surface as 400, not 500. Verifies the controller passes file=None to AudioService.transcript_asr @@ -992,7 +1044,7 @@ class TestTrialChatAudioApi: assert exc_info.value.code == 400 - def test_audio_too_large(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_audio_too_large(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1015,7 +1067,7 @@ class TestTrialChatAudioApi: trial_app_chat, ) - def test_unsupported_audio_type(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_unsupported_audio_type(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1038,7 +1090,7 @@ class TestTrialChatAudioApi: trial_app_chat, ) - def test_provider_not_support_tts(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_provider_not_support_tts(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1061,7 +1113,7 @@ class TestTrialChatAudioApi: trial_app_chat, ) - def test_speech_to_text_disabled(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_speech_to_text_disabled(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) file_data = _file_data() @@ -1083,7 +1135,7 @@ class TestTrialChatAudioApi: trial_app_chat, ) - def test_provider_not_init(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_provider_not_init(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1102,7 +1154,7 @@ class TestTrialChatAudioApi: trial_app_chat, ) - def test_quota_exceeded(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_quota_exceeded(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1123,22 +1175,34 @@ class TestTrialChatAudioApi: class TestTrialChatTextApi: - def test_success(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_success( + self, + app: Flask, + trial_app_chat: App, + account: Account, + trial_app_usage: MagicMock, + ) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) with ( app.test_request_context("/", json={"text": "hello", "voice": "en-US"}), patch.object(module.AudioService, "transcript_tts", return_value={"audio": "base64_data"}), - patch.object(module.RecommendedAppService, "add_trial_app_record"), ): result = method( api, TextToSpeechRequest.model_validate(request.get_json(silent=True) or {}), account, trial_app_chat ) assert result == {"audio": "base64_data"} + trial_app_usage.record.assert_called_once_with(app_id="a-chat", account_id="u1") - def test_success_with_message_ref(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_success_with_message_ref( + self, + app: Flask, + trial_app_chat: App, + account: Account, + trial_app_usage: MagicMock, + ) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) transcript_tts = MagicMock(return_value={"audio": "base64_data"}) @@ -1147,7 +1211,6 @@ class TestTrialChatTextApi: with ( app.test_request_context("/", json={"text": "hello", "message_id": "message-1"}), patch.object(module.AudioService, "transcript_tts", transcript_tts), - patch.object(module.RecommendedAppService, "add_trial_app_record"), ): result = method( api, TextToSpeechRequest.model_validate(request.get_json(silent=True) or {}), account, trial_app_chat @@ -1159,8 +1222,9 @@ class TestTrialChatTextApi: "message-1", account_id="u1", ) + trial_app_usage.record.assert_called_once_with(app_id="a-chat", account_id="u1") - def test_app_config_broken(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_app_config_broken(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1180,7 +1244,7 @@ class TestTrialChatTextApi: trial_app_chat, ) - def test_provider_not_support(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_provider_not_support(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1200,7 +1264,7 @@ class TestTrialChatTextApi: trial_app_chat, ) - def test_audio_too_large(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_audio_too_large(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1220,7 +1284,7 @@ class TestTrialChatTextApi: trial_app_chat, ) - def test_no_audio_uploaded(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_no_audio_uploaded(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1240,7 +1304,7 @@ class TestTrialChatTextApi: trial_app_chat, ) - def test_provider_not_init(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_provider_not_init(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1256,7 +1320,7 @@ class TestTrialChatTextApi: trial_app_chat, ) - def test_quota_exceeded(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_quota_exceeded(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1272,7 +1336,7 @@ class TestTrialChatTextApi: trial_app_chat, ) - def test_model_not_support(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_model_not_support(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1288,7 +1352,7 @@ class TestTrialChatTextApi: trial_app_chat, ) - def test_invoke_error(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_invoke_error(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1306,14 +1370,14 @@ class TestTrialChatTextApi: class TestTrialAppWorkflowTaskStopApi: - def test_not_workflow_app(self, app: Flask, trial_app_chat: MagicMock) -> None: + def test_not_workflow_app(self, app: Flask, trial_app_chat: App) -> None: api = module.TrialAppWorkflowTaskStopApi() with app.test_request_context("/", json={"inputs": {}}): with pytest.raises(NotWorkflowAppError): api.post(trial_app_chat, str(uuid4())) - def test_success(self, app: Flask, trial_app_workflow: MagicMock) -> None: + def test_success(self, app: Flask, trial_app_workflow: App) -> None: api = module.TrialAppWorkflowTaskStopApi() task_id = str(uuid4()) @@ -1330,7 +1394,6 @@ class TestTrialAppWorkflowTaskStopApi: class TestTrialSitApi: - @pytest.mark.parametrize("sqlite_session", [(Site,)], indirect=True) def test_no_site( self, app: Flask, @@ -1338,14 +1401,12 @@ class TestTrialSitApi: ) -> None: api = module.TrialSitApi() method = unwrap(api.get) - app_model = MagicMock() - app_model.id = str(uuid4()) + app_model = _app(app_id=str(uuid4()), mode=AppMode.CHAT) with app.test_request_context("/"): with pytest.raises(Forbidden): method(api, sqlite_session, app_model) - @pytest.mark.parametrize("sqlite_session", [(Site,)], indirect=True) def test_archived_tenant( self, app: Flask, @@ -1354,8 +1415,9 @@ class TestTrialSitApi: api = module.TrialSitApi() method = unwrap(api.get) - app_model = SimpleNamespace(id=str(uuid4()), tenant_id="tenant-1") - tenant = SimpleNamespace(status=TenantStatus.ARCHIVE) + app_model = _app(app_id=str(uuid4()), mode=AppMode.CHAT) + tenant = Tenant(name="Archived Tenant", status=TenantStatus.ARCHIVE) + tenant.id = app_model.tenant_id _persist_site(sqlite_session, app_model.id) with ( @@ -1367,7 +1429,6 @@ class TestTrialSitApi: get_tenant_by_id.assert_called_once_with("tenant-1", session=sqlite_session) - @pytest.mark.parametrize("sqlite_session", [(Site,)], indirect=True) def test_success( self, app: Flask, @@ -1376,8 +1437,9 @@ class TestTrialSitApi: api = module.TrialSitApi() method = unwrap(api.get) - app_model = SimpleNamespace(id=str(uuid4()), tenant_id="tenant-1") - tenant = SimpleNamespace(status=TenantStatus.NORMAL) + app_model = _app(app_id=str(uuid4()), mode=AppMode.CHAT) + tenant = Tenant(name="Active Tenant", status=TenantStatus.NORMAL) + tenant.id = app_model.tenant_id site = _persist_site(sqlite_session, app_model.id) with ( @@ -1396,56 +1458,69 @@ class TestTrialSitApi: class TestAppWorkflowApi: - def test_uses_injected_session(self, unbound_session: Session) -> None: + def test_uses_injected_session(self, sqlite_session: Session) -> None: api = module.AppWorkflowApi() method = unwrap(api.get) - created_by = SimpleNamespace(id="account-1", name="Creator", email="creator@example.com") - workflow = SimpleNamespace( - id="workflow-1", - graph_dict={"nodes": []}, - features_dict={}, - unique_hash="workflow-hash", - version="draft", - marked_name="", - marked_comment="", - created_at=datetime(2024, 1, 1, tzinfo=UTC), - updated_at=datetime(2024, 1, 2, tzinfo=UTC), - environment_variables=[ - SecretVariable( - id="env-secret", - name="api_key", - value="plaintext-secret", - ), - LLMEnvironmentVariable( - id="env-llm", - name="shared_model", - value={"provider": "provider", "name": "model", "mode": "chat"}, - ), - ], - conversation_variables=[ - StringVariable( - id="conversation-variable-1", - name="topic", - value="sqlite", - selector=["conversation", "topic"], - ) - ], - rag_pipeline_variables=[], - get_created_by_account=MagicMock(return_value=created_by), - get_updated_by_account=MagicMock(return_value=None), - get_tool_published=MagicMock(return_value=True), + created_by = Account(name="Creator", email="creator@example.com") + created_by.id = "account-1" + app_model = _app(app_id="app-1", mode=AppMode.WORKFLOW) + with patch("models.workflow.encrypter.encrypt_token", return_value="encrypted-secret"): + workflow = Workflow.new( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + type="workflow", + version="draft", + graph='{"nodes": []}', + features="{}", + created_by=created_by.id, + environment_variables=[ + SecretVariable( + id="env-secret", + name="api_key", + value="plaintext-secret", + ), + LLMEnvironmentVariable( + id="env-llm", + name="shared_model", + value={"provider": "provider", "name": "model", "mode": "chat"}, + ), + ], + conversation_variables=[ + StringVariable( + id="conversation-variable-1", + name="topic", + value="sqlite", + selector=["conversation", "topic"], + ) + ], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-1" + workflow.created_at = datetime(2024, 1, 1, tzinfo=UTC) + workflow.updated_at = datetime(2024, 1, 2, tzinfo=UTC) + app_model.workflow_id = workflow.id + tool_provider = WorkflowToolProvider( + name="trial-workflow", + label="Trial Workflow", + icon="icon", + app_id=app_model.id, + version="1.0.0", + user_id=created_by.id, + tenant_id=app_model.tenant_id, + description="Trial workflow provider", + parameter_configuration="[]", ) - app_model = SimpleNamespace( - workflow_id="workflow-1", - workflow_with_session=MagicMock(return_value=workflow), - ) - result = method(api, unbound_session, app_model) + sqlite_session.add_all([created_by, app_model, workflow, tool_provider]) + sqlite_session.commit() + + with patch("models.workflow.encrypter.decrypt_token", return_value="plaintext-secret"): + result = method(api, sqlite_session, app_model) assert result == { "id": "workflow-1", "graph": {"nodes": []}, "features": {}, - "hash": "workflow-hash", + "hash": workflow.unique_hash, "version": "draft", "marked_name": "", "marked_comment": "", @@ -1461,7 +1536,7 @@ class TestAppWorkflowApi: "id": "env-secret", "name": "api_key", "description": "", - "selector": [], + "selector": ["env", "api_key"], }, { "value_type": "llm", @@ -1469,7 +1544,7 @@ class TestAppWorkflowApi: "id": "env-llm", "name": "shared_model", "description": "", - "selector": [], + "selector": ["env", "shared_model"], }, ], "conversation_variables": [ @@ -1483,14 +1558,10 @@ class TestAppWorkflowApi: ], "rag_pipeline_variables": [], } - app_model.workflow_with_session.assert_called_once_with(session=unbound_session) - workflow.get_created_by_account.assert_called_once_with(session=unbound_session) - workflow.get_updated_by_account.assert_called_once_with(session=unbound_session) - workflow.get_tool_published.assert_called_once_with(session=unbound_session) class TestTrialChatAudioApiExceptionHandlers: - def test_provider_not_init(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_provider_not_init(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1513,7 +1584,7 @@ class TestTrialChatAudioApiExceptionHandlers: trial_app_chat, ) - def test_quota_exceeded(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_quota_exceeded(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1536,7 +1607,7 @@ class TestTrialChatAudioApiExceptionHandlers: trial_app_chat, ) - def test_invoke_error(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_invoke_error(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatAudioApi() method = unwrap(api.post) @@ -1561,7 +1632,7 @@ class TestTrialChatAudioApiExceptionHandlers: class TestTrialChatTextApiExceptionHandlers: - def test_app_config_broken(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_app_config_broken(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) @@ -1581,7 +1652,7 @@ class TestTrialChatTextApiExceptionHandlers: trial_app_chat, ) - def test_unsupported_audio_type(self, app: Flask, trial_app_chat: MagicMock, account: Account) -> None: + def test_unsupported_audio_type(self, app: Flask, trial_app_chat: App, account: Account) -> None: api = module.TrialChatTextApi() method = unwrap(api.post) diff --git a/api/tests/unit_tests/controllers/console/explore/test_workflow.py b/api/tests/unit_tests/controllers/console/explore/test_workflow.py index 247669fe125..835b316dab7 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/explore/test_workflow.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch import pytest from flask import Flask +from sqlalchemy.orm import Session from werkzeug.exceptions import InternalServerError from controllers.common.controller_schemas import WorkflowRunPayload @@ -12,7 +13,8 @@ from controllers.console.explore.workflow import ( InstalledAppWorkflowTaskStopApi, ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError -from models.model import AppMode +from models import Account +from models.model import App, AppMode, InstalledApp from services.errors.llm import InvokeRateLimitError @@ -23,31 +25,26 @@ def app(): return app -@pytest.fixture -def user(): - return MagicMock() - - -@pytest.fixture -def workflow_app(): - app = MagicMock() - app.mode = AppMode.WORKFLOW - return app - - -@pytest.fixture -def installed_workflow_app(workflow_app): - installed_app = MagicMock(app=workflow_app) - installed_app.app_with_session.return_value = workflow_app - return installed_app - - -@pytest.fixture -def non_workflow_installed_app(): - app = MagicMock() - app.mode = AppMode.CHAT - installed_app = MagicMock(app=app) - installed_app.app_with_session.return_value = app +def make_installed_app(session: Session, *, mode: AppMode) -> InstalledApp: + app = App( + tenant_id="owner-tenant", + name="Explore App", + mode=mode, + enable_site=True, + enable_api=False, + ) + session.add(app) + session.flush() + installed_app = InstalledApp( + tenant_id="viewer-tenant", + app_id=app.id, + app_owner_tenant_id=app.tenant_id, + position=0, + is_pinned=False, + last_used_at=None, + ) + session.add(installed_app) + session.commit() return installed_app @@ -57,24 +54,28 @@ def payload(): class TestInstalledAppWorkflowRunApi: - def test_not_workflow_app(self, app: Flask, non_workflow_installed_app): + def test_not_workflow_app(self, app: Flask, sqlite_session: Session): api = InstalledAppWorkflowRunApi() method = unwrap(api.post) + installed_app = make_installed_app(sqlite_session, mode=AppMode.CHAT) + user = Account(name="User", email="user@example.com") with app.test_request_context("/"): with pytest.raises(NotWorkflowAppError): method( api, WorkflowRunPayload.model_validate({"inputs": {}}), - MagicMock(), - MagicMock(), - non_workflow_installed_app, + sqlite_session, + user, + installed_app, ) - def test_success(self, app: Flask, installed_workflow_app, user, payload): + def test_success(self, app: Flask, sqlite_session: Session, payload): api = InstalledAppWorkflowRunApi() method = unwrap(api.post) req_data = WorkflowRunPayload.model_validate(payload) + installed_app = make_installed_app(sqlite_session, mode=AppMode.WORKFLOW) + user = Account(name="User", email="user@example.com") with ( app.test_request_context("/", json=payload), @@ -83,16 +84,18 @@ class TestInstalledAppWorkflowRunApi: return_value=MagicMock(), ) as generate_mock, ): - result = method(api, req_data, MagicMock(), user, installed_workflow_app) + result = method(api, req_data, sqlite_session, user, installed_app) generate_mock.assert_called_once() assert generate_mock.call_args.kwargs["user"] is user assert result is not None - def test_rate_limit_error(self, app: Flask, installed_workflow_app, user, payload): + def test_rate_limit_error(self, app: Flask, sqlite_session: Session, payload): api = InstalledAppWorkflowRunApi() method = unwrap(api.post) req_data = WorkflowRunPayload.model_validate(payload) + installed_app = make_installed_app(sqlite_session, mode=AppMode.WORKFLOW) + user = Account(name="User", email="user@example.com") with ( app.test_request_context("/", json=payload), @@ -102,12 +105,14 @@ class TestInstalledAppWorkflowRunApi: ), ): with pytest.raises(InvokeRateLimitHttpError): - method(api, req_data, MagicMock(), user, installed_workflow_app) + method(api, req_data, sqlite_session, user, installed_app) - def test_unexpected_exception(self, app: Flask, installed_workflow_app, user, payload): + def test_unexpected_exception(self, app: Flask, sqlite_session: Session, payload): api = InstalledAppWorkflowRunApi() method = unwrap(api.post) req_data = WorkflowRunPayload.model_validate(payload) + installed_app = make_installed_app(sqlite_session, mode=AppMode.WORKFLOW) + user = Account(name="User", email="user@example.com") with ( app.test_request_context("/", json=payload), @@ -117,26 +122,28 @@ class TestInstalledAppWorkflowRunApi: ), ): with pytest.raises(InternalServerError): - method(api, req_data, MagicMock(), user, installed_workflow_app) + method(api, req_data, sqlite_session, user, installed_app) class TestInstalledAppWorkflowTaskStopApi: - def test_not_workflow_app(self, non_workflow_installed_app): + def test_not_workflow_app(self, sqlite_session: Session): api = InstalledAppWorkflowTaskStopApi() method = unwrap(api.post) + installed_app = make_installed_app(sqlite_session, mode=AppMode.CHAT) with pytest.raises(NotWorkflowAppError): - method(api, MagicMock(), non_workflow_installed_app, "task-1") + method(api, sqlite_session, installed_app, "task-1") - def test_success(self, installed_workflow_app): + def test_success(self, sqlite_session: Session): api = InstalledAppWorkflowTaskStopApi() method = unwrap(api.post) + installed_app = make_installed_app(sqlite_session, mode=AppMode.WORKFLOW) with ( patch("controllers.console.explore.workflow.AppQueueManager.set_stop_flag_no_user_check") as stop_flag, patch("controllers.console.explore.workflow.GraphEngineManager.send_stop_command") as send_stop, ): - result = method(api, MagicMock(), installed_workflow_app, "task-1") + result = method(api, sqlite_session, installed_app, "task-1") stop_flag.assert_called_once_with("task-1") send_stop.assert_called_once_with("task-1") diff --git a/api/tests/unit_tests/controllers/console/explore/test_wraps.py b/api/tests/unit_tests/controllers/console/explore/test_wraps.py index a1da7916f02..fee341593d7 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/explore/test_wraps.py @@ -1,15 +1,16 @@ from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy.orm import Session, scoped_session -from werkzeug.exceptions import Forbidden, NotFound +from werkzeug.exceptions import NotFound import controllers.console.explore.wraps as wraps_module import models.model as model_module from controllers.console.explore.error import ( AppAccessDeniedError, + TrialAppFeatureDisabledError, TrialAppLimitExceeded, TrialAppNotAllowed, ) @@ -260,23 +261,27 @@ def test_trial_feature_enable_disabled(): def view(): return "ok" - with patch( - "controllers.console.explore.wraps.RecommendedAppService.is_trial_app_enabled", - return_value=False, - ): - with pytest.raises(Forbidden): + services = MagicMock() + services.recommended_app_queries.is_trial_enabled.return_value = False + with patch("controllers.console.explore.wraps.application_services", return_value=services): + with pytest.raises(TrialAppFeatureDisabledError) as exc_info: view() + assert exc_info.value.data == { + "code": "trial_app_feature_disabled", + "message": "Trial app feature is not enabled.", + "status": 403, + } + def test_trial_feature_enable_enabled(): @trial_feature_enable def view(): return "ok" - with patch( - "controllers.console.explore.wraps.RecommendedAppService.is_trial_app_enabled", - return_value=True, - ): + services = MagicMock() + services.recommended_app_queries.is_trial_enabled.return_value = True + with patch("controllers.console.explore.wraps.application_services", return_value=services): assert view() == "ok" diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py index b5b2d79411b..ad4de5468ab 100644 --- a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py +++ b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py @@ -197,6 +197,33 @@ def test_published_workflow_post_returns_400_when_publish_fails( assert snippet.name == "Snippet" +@pytest.mark.parametrize("sqlite_session", [(CustomizedSnippet,)], indirect=True) +def test_published_workflow_post_returns_success( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, +) -> None: + user = _account("account-1") + snippet = _snippet() + sqlite_session.add(snippet) + sqlite_session.commit() + workflow = SimpleNamespace(created_at=datetime(2026, 8, 17, 12, 0, 0)) + monkeypatch.setattr(snippet_workflow_module, "db", SimpleNamespace(engine=sqlite_engine)) + monkeypatch.setattr( + snippet_workflow_module, + "_snippet_service", + lambda: SimpleNamespace(publish_workflow=Mock(return_value=workflow)), + ) + + api = snippet_workflow_module.SnippetPublishedWorkflowApi() + handler = unwrap(api.post) + with app.test_request_context("/snippets/snippet-1/workflows/publish", method="POST", json={}): + response = handler(api, user, snippet) + + assert response["result"] == "success" + + def test_default_block_configs_delegates_to_service(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: get_default_block_configs = Mock(return_value=[{"type": "llm"}]) monkeypatch.setattr( @@ -505,6 +532,70 @@ def test_update_published_snippet_workflow_raises_not_found( assert snippet.name == "Snippet" +def test_delete_published_snippet_workflow_succeeds(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + snippet = _snippet() + delete_workflow = Mock(return_value=True) + monkeypatch.setattr( + snippet_workflow_module, + "SnippetService", + lambda: SimpleNamespace(delete_workflow=delete_workflow), + ) + + api = snippet_workflow_module.SnippetWorkflowByIdApi() + handler = unwrap(api.delete) + + with app.test_request_context("/snippets/snippet-1/workflows/workflow-1", method="DELETE"): + response, status_code = handler(api, snippet, workflow_id="workflow-1") + + assert status_code == 204 + assert response is None + delete_workflow.assert_called_once() + delete_call = delete_workflow.call_args.kwargs + assert isinstance(delete_call["session"], Session) + assert delete_call["snippet"] is snippet + assert delete_call["workflow_id"] == "workflow-1" + + +def test_delete_published_snippet_workflow_raises_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + def delete_missing_workflow(**_kwargs): + raise ValueError("Workflow with ID missing-workflow not found") + + monkeypatch.setattr( + snippet_workflow_module, + "SnippetService", + lambda: SimpleNamespace(delete_workflow=Mock(side_effect=delete_missing_workflow)), + ) + + api = snippet_workflow_module.SnippetWorkflowByIdApi() + handler = unwrap(api.delete) + + with app.test_request_context("/snippets/snippet-1/workflows/missing-workflow", method="DELETE"): + with pytest.raises(NotFound): + handler(api, _snippet(), workflow_id="missing-workflow") + + +def test_delete_published_snippet_workflow_raises_bad_request_when_in_use( + app: Flask, monkeypatch: pytest.MonkeyPatch +) -> None: + def delete_active_workflow(**_kwargs): + raise snippet_workflow_module.WorkflowInUseError("Cannot delete workflow that is currently in use") + + monkeypatch.setattr( + snippet_workflow_module, + "SnippetService", + lambda: SimpleNamespace(delete_workflow=Mock(side_effect=delete_active_workflow)), + ) + + api = snippet_workflow_module.SnippetWorkflowByIdApi() + handler = unwrap(api.delete) + + with app.test_request_context("/snippets/snippet-1/workflows/workflow-1", method="DELETE"): + with pytest.raises(HTTPException) as exc_info: + handler(api, _snippet(), workflow_id="workflow-1") + + assert exc_info.value.code == 400 + + def test_workflow_run_detail_raises_not_found_when_run_missing(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: snippet = _snippet() monkeypatch.setattr( diff --git a/api/tests/unit_tests/controllers/console/tag/test_tags.py b/api/tests/unit_tests/controllers/console/tag/test_tags.py index b5d383e00e6..132dd6fff0c 100644 --- a/api/tests/unit_tests/controllers/console/tag/test_tags.py +++ b/api/tests/unit_tests/controllers/console/tag/test_tags.py @@ -1,12 +1,9 @@ -from collections.abc import Iterator from types import SimpleNamespace -from unittest.mock import PropertyMock, patch +from unittest.mock import MagicMock, patch import pytest from flask import Flask -from sqlalchemy import Engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker -from werkzeug.exceptions import Forbidden +from werkzeug.exceptions import Forbidden, NotFound, UnprocessableEntity import controllers.console.tag.tags as module from controllers.console import console_ns @@ -21,194 +18,146 @@ from controllers.console.tag.tags import ( TagUpdateDeleteApi, TagUpdateRequestPayload, ) +from machinery.context import RequestContext from models import Account from models.account import AccountStatus, TenantAccountRole -from models.base import TypeBase from models.enums import TagType -from models.model import Tag -from services.tag_service import UpdateTagPayload +from services.tag_application_service import ( + TagApplicationError, + TagBindingInput, + TagBindingTargetNotFoundError, + TagNameConflictError, + TagNotFoundError, + TagSummary, + UpdateTagInput, +) def unwrap(func): - """ - Recursively unwrap decorated functions. - """ while hasattr(func, "__wrapped__"): func = func.__wrapped__ return func @pytest.fixture -def app(): +def app() -> Flask: app = Flask("test_tag") app.config["TESTING"] = True return app -@pytest.fixture(autouse=True) -def sqlite_db_session( - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, -) -> Iterator[scoped_session[Session]]: - TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[Tag.__tablename__]]) - session_registry = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - monkeypatch.setattr(module.db, "session", session_registry) - try: - yield session_registry - finally: - session_registry.remove() - - -def _assert_sqlite_session(session: object, sqlite_engine: Engine) -> None: - assert isinstance(session, Session) - assert session.get_bind() is sqlite_engine - assert session.is_active - - @pytest.fixture -def admin_user(): +def request_context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id=None, + account_id="user-1", + active_workspace_id="tenant-1", + ) + + +def _account(role: TenantAccountRole) -> Account: account = Account( - name="Admin User", - email="admin@example.com", + name="Tag User", + email=f"{role.value}@example.com", status=AccountStatus.ACTIVE, ) account.id = "user-1" - account.role = TenantAccountRole.OWNER + account.role = role return account @pytest.fixture -def readonly_user(): - account = Account( - name="Readonly User", - email="readonly@example.com", - status=AccountStatus.ACTIVE, - ) - account.id = "user-2" - account.role = TenantAccountRole.NORMAL - return account - - -@pytest.fixture -def tag(sqlite_db_session: scoped_session[Session]): - tag = Tag( - tenant_id="tenant-1", - name="test-tag", - type=TagType.KNOWLEDGE, - created_by="user-1", - ) - tag.id = "tag-1" - sqlite_db_session.add(tag) - sqlite_db_session.commit() - return tag - - -@pytest.fixture -def payload_patch(): - def _patch(payload): - return patch.object( - type(console_ns), - "payload", - new_callable=PropertyMock, - return_value=payload, - ) - - return _patch +def tags_service() -> MagicMock: + tags = MagicMock() + with patch.object(module, "application_services", return_value=SimpleNamespace(tags=tags)): + yield tags class TestTagListApi: - def test_get_success(self, app: Flask): - api = TagListApi() - method = unwrap(api.get) + @pytest.mark.parametrize("url", ["/", "/?type="]) + def test_get_requires_non_empty_type(self, app: Flask, url: str) -> None: + class Handler: + @module.model_validate(TagListQueryParam) + def get(self, req_data: TagListQueryParam) -> TagListQueryParam: + return req_data + + with app.test_request_context(url, method="GET"): + with pytest.raises(UnprocessableEntity): + Handler().get() + + def test_get_uses_application_service( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.list_tags.return_value = (TagSummary("tag-1", "Tag", "knowledge", 2),) with app.test_request_context("/?type=knowledge"): - with ( - patch( - "controllers.console.tag.tags.TagService.get_tags", - return_value=[ - SimpleNamespace( - id="1", - name="tag", - type=TagType.KNOWLEDGE, - binding_count=1, - ) - ], - ), - ): - result, status = method(api, TagListQueryParam(type="knowledge"), "tenant-1") + result, status = unwrap(TagListApi().get)( + TagListApi(), + TagListQueryParam(type="knowledge"), + request_context, + ) + tags_service.list_tags.assert_called_once_with(request_context, "knowledge", None) assert status == 200 - assert result == [{"id": "1", "name": "tag", "type": "knowledge", "binding_count": "1"}] + assert result == [{"id": "tag-1", "name": "Tag", "type": "knowledge", "binding_count": "2"}] - def test_get_snippet_tags(self, app: Flask, sqlite_engine: Engine): - api = TagListApi() - method = unwrap(api.get) + def test_get_snippet_tags_uses_same_query_boundary( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.list_tags.return_value = (TagSummary("tag-1", "Snippet", "snippet", 1),) with app.test_request_context("/?type=snippet"): - with ( - patch( - "controllers.console.tag.tags.TagService.get_tags", - return_value=[ - SimpleNamespace( - id="1", - name="snippet-tag", - type=TagType.SNIPPET, - binding_count=1, - ) - ], - ) as get_tags_mock, - ): - result, status = method(api, TagListQueryParam(type="snippet"), "tenant-1") + result, status = unwrap(TagListApi().get)( + TagListApi(), + TagListQueryParam(type="snippet"), + request_context, + ) - get_tags_mock.assert_called_once() - assert get_tags_mock.call_args.args == ("snippet", "tenant-1", None) - _assert_sqlite_session(get_tags_mock.call_args.kwargs["session"], sqlite_engine) + tags_service.list_tags.assert_called_once_with(request_context, "snippet", None) assert status == 200 - assert result == [{"id": "1", "name": "snippet-tag", "type": "snippet", "binding_count": "1"}] + assert result[0]["type"] == "snippet" - def test_post_success(self, app: Flask, admin_user, tag): - api = TagListApi() - method = unwrap(api.post) + def test_post_preserves_dataset_editor_permission( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.create_tag.return_value = TagSummary("tag-1", "Tag", "knowledge", 0) + dataset_operator = _account(TenantAccountRole.DATASET_OPERATOR) - payload = {"name": "test-tag", "type": "knowledge"} - req_data = TagBasePayload.model_validate(payload) - - with app.test_request_context("/", json=payload): - with ( - patch( - "controllers.console.tag.tags.TagService.save_tags", - return_value=tag, - ), - ): - result, status = method(api, req_data, admin_user) + with ( + app.test_request_context("/", json={"name": "Tag", "type": "knowledge"}), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(dataset_operator, "tenant-1")), + ): + result, status = unwrap(TagListApi().post)( + TagListApi(), + TagBasePayload(name="Tag", type=TagType.KNOWLEDGE), + request_context, + ) assert status == 200 - assert result["name"] == "test-tag" assert result["binding_count"] == "0" + tags_service.create_tag.assert_called_once() - def test_post_snippet_tag_checks_snippet_rbac_when_enabled(self, app: Flask, admin_user, tag): - api = TagListApi() - method = unwrap(api.post) + def test_post_snippet_tag_checks_rbac( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.create_tag.return_value = TagSummary("tag-1", "Snippet", "snippet", 0) + owner = _account(TenantAccountRole.OWNER) - payload = {"name": "snippet-tag", "type": "snippet"} - req_data = TagBasePayload.model_validate(payload) + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", True), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + patch.object(module, "enforce_rbac_access") as enforce_rbac_access, + ): + unwrap(TagListApi().post)( + TagListApi(), + TagBasePayload(name="Snippet", type=TagType.SNIPPET), + request_context, + ) - with app.test_request_context("/", json=payload): - with ( - patch("controllers.console.tag.tags.dify_config.RBAC_ENABLED", True), - patch( - "controllers.console.tag.tags.current_account_with_tenant", - return_value=(admin_user, "tenant-1"), - ), - patch("controllers.console.tag.tags.enforce_rbac_access") as enforce_mock, - patch( - "controllers.console.tag.tags.TagService.save_tags", - return_value=tag, - ), - ): - method(api, req_data, admin_user) - - enforce_mock.assert_called_once_with( + enforce_rbac_access.assert_called_once_with( tenant_id="tenant-1", account_id="user-1", resource_type=module.RBACResourceScope.WORKSPACE, @@ -216,256 +165,305 @@ class TestTagListApi: resource_required=False, ) - def test_post_forbidden(self, app: Flask, readonly_user): - api = TagListApi() - method = unwrap(api.post) + def test_post_rejects_read_only_member(self, app: Flask, request_context: RequestContext) -> None: + readonly = _account(TenantAccountRole.NORMAL) - with app.test_request_context("/"): + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(readonly, "tenant-1")), + ): with pytest.raises(Forbidden): - method(api, TagBasePayload(name="test", type=TagType.KNOWLEDGE), readonly_user) + unwrap(TagListApi().post)( + TagListApi(), + TagBasePayload(name="Tag", type=TagType.KNOWLEDGE), + request_context, + ) + + def test_post_maps_name_conflict_to_legacy_value_error( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.create_tag.side_effect = TagNameConflictError() + owner = _account(TenantAccountRole.OWNER) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + with pytest.raises(ValueError, match="Tag name already exists") as exc_info: + unwrap(TagListApi().post)( + TagListApi(), + TagBasePayload(name="Tag", type=TagType.KNOWLEDGE), + request_context, + ) + + assert exc_info.value.__cause__ is None + assert exc_info.value.__suppress_context__ is True + + def test_post_does_not_coerce_unknown_application_error_to_transport_error( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.create_tag.side_effect = TagApplicationError("unexpected") + owner = _account(TenantAccountRole.OWNER) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + with pytest.raises(TagApplicationError, match="unexpected"): + unwrap(TagListApi().post)( + TagListApi(), + TagBasePayload(name="Tag", type=TagType.KNOWLEDGE), + request_context, + ) class TestTagUpdateDeleteApi: - def test_patch_success(self, app: Flask, admin_user, tag, sqlite_engine: Engine): - api = TagUpdateDeleteApi() - method = unwrap(api.patch) - - payload = {"name": "updated"} - req_data = TagUpdateRequestPayload.model_validate(payload) - - with app.test_request_context("/", json=payload): - with ( - patch( - "controllers.console.tag.tags.TagService.update_tags", - return_value=tag, - ) as update_tags_mock, - patch( - "controllers.console.tag.tags.TagService.get_tag_binding_count", - return_value=3, - ), - ): - result, status = method(api, req_data, admin_user, "tag-1") - - assert status == 200 - update_payload, tag_id, session = update_tags_mock.call_args.args - assert update_payload == UpdateTagPayload(name="updated") - assert tag_id == "tag-1" - _assert_sqlite_session(session, sqlite_engine) - assert result["binding_count"] == "3" - - def test_patch_forbidden(self, app: Flask, readonly_user): - api = TagUpdateDeleteApi() - method = unwrap(api.patch) - - with app.test_request_context("/"): - with pytest.raises(Forbidden): - method(api, TagUpdateRequestPayload(name="test"), readonly_user, "tag-1") - - def test_delete_success(self, app: Flask, admin_user, sqlite_engine: Engine): - api = TagUpdateDeleteApi() - method = unwrap(api.delete) + def test_patch_authorizes_snippet_before_update( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.get_tag_type.return_value = "snippet" + tags_service.update_tag.return_value = TagSummary("tag-1", "Updated", "snippet", 3) + owner = _account(TenantAccountRole.OWNER) with ( app.test_request_context("/"), - patch("controllers.console.tag.tags.TagService.delete_tag") as delete_mock, + patch.object(module.dify_config, "RBAC_ENABLED", True), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + patch.object(module, "enforce_rbac_access") as enforce_rbac_access, ): - result, status = method(api, "tag-1") + result, status = unwrap(TagUpdateDeleteApi().patch)( + TagUpdateDeleteApi(), + TagUpdateRequestPayload(name="Updated"), + request_context, + "tag-1", + ) - delete_mock.assert_called_once() - tag_id, session = delete_mock.call_args.args - assert tag_id == "tag-1" - _assert_sqlite_session(session, sqlite_engine) - assert status == 204 - - def test_delete_snippet_tag_checks_type_in_current_tenant( - self, - app: Flask, - admin_user, - sqlite_db_session: scoped_session[Session], - sqlite_engine: Engine, - ): - api = TagUpdateDeleteApi() - method = unwrap(api.delete) - tag = Tag( - tenant_id="tenant-1", - name="snippet-tag", - type=TagType.SNIPPET, - created_by="user-1", - ) - tag.id = "tag-1" - sqlite_db_session.add(tag) - sqlite_db_session.commit() - - with ( - app.test_request_context("/"), - patch("controllers.console.tag.tags.dify_config.RBAC_ENABLED", True), - patch( - "controllers.console.tag.tags.current_account_with_tenant", - return_value=(admin_user, "tenant-1"), - ), - patch("controllers.console.tag.tags.enforce_rbac_access") as enforce_mock, - patch("controllers.console.tag.tags.TagService.delete_tag") as delete_mock, - ): - result, status = method(api, "tag-1") - - enforce_mock.assert_called_once_with( + enforce_rbac_access.assert_called_once_with( tenant_id="tenant-1", account_id="user-1", resource_type=module.RBACResourceScope.WORKSPACE, scene=module.RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False, ) - delete_mock.assert_called_once() - tag_id, session = delete_mock.call_args.args - assert tag_id == "tag-1" - _assert_sqlite_session(session, sqlite_engine) - assert result == "" - assert status == 204 + tags_service.update_tag.assert_called_once_with(request_context, "tag-1", UpdateTagInput(name="Updated")) + assert status == 200 + assert result["binding_count"] == "3" - def test_delete_does_not_apply_snippet_rbac_to_tag_from_another_tenant( - self, - app: Flask, - admin_user, - sqlite_db_session: scoped_session[Session], - sqlite_engine: Engine, - ): - api = TagUpdateDeleteApi() - method = unwrap(api.delete) - tag = Tag( - tenant_id="other-tenant", - name="other-tenant-snippet-tag", - type=TagType.SNIPPET, - created_by="other-user", - ) - tag.id = "tag-1" - sqlite_db_session.add(tag) - sqlite_db_session.commit() + def test_patch_rejects_read_only_member(self, app: Flask, request_context: RequestContext) -> None: + readonly = _account(TenantAccountRole.NORMAL) with ( app.test_request_context("/"), - patch("controllers.console.tag.tags.dify_config.RBAC_ENABLED", True), - patch( - "controllers.console.tag.tags.current_account_with_tenant", - return_value=(admin_user, "tenant-1"), - ), - patch("controllers.console.tag.tags.enforce_rbac_access") as enforce_mock, - patch("controllers.console.tag.tags.TagService.delete_tag") as delete_mock, + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(readonly, "tenant-1")), ): - result, status = method(api, "tag-1") + with pytest.raises(Forbidden): + unwrap(TagUpdateDeleteApi().patch)( + TagUpdateDeleteApi(), + TagUpdateRequestPayload(name="Updated"), + request_context, + "tag-1", + ) - enforce_mock.assert_not_called() - delete_mock.assert_called_once() - tag_id, session = delete_mock.call_args.args - assert tag_id == "tag-1" - _assert_sqlite_session(session, sqlite_engine) - assert result == "" - assert status == 204 + def test_patch_maps_missing_tag_to_not_found( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.update_tag.side_effect = TagNotFoundError() + owner = _account(TenantAccountRole.OWNER) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + with pytest.raises(NotFound, match="Tag not found") as exc_info: + unwrap(TagUpdateDeleteApi().patch)( + TagUpdateDeleteApi(), + TagUpdateRequestPayload(name="Updated"), + request_context, + "tag-1", + ) + + assert exc_info.value.__cause__ is None + assert exc_info.value.__suppress_context__ is True + + def test_delete_does_not_grant_dataset_operator_legacy_edit_permission( + self, app: Flask, request_context: RequestContext + ) -> None: + dataset_operator = _account(TenantAccountRole.DATASET_OPERATOR) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(dataset_operator, "tenant-1")), + ): + with pytest.raises(Forbidden): + unwrap(TagUpdateDeleteApi().delete)(TagUpdateDeleteApi(), request_context, "tag-1") + + def test_delete_calls_application_service( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + owner = _account(TenantAccountRole.OWNER) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + result, status = unwrap(TagUpdateDeleteApi().delete)(TagUpdateDeleteApi(), request_context, "tag-1") + + tags_service.delete_tag.assert_called_once_with(request_context, "tag-1") + assert (result, status) == ("", 204) + + def test_delete_snippet_tag_checks_type_in_current_workspace( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.get_tag_type.return_value = "snippet" + owner = _account(TenantAccountRole.OWNER) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", True), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + patch.object(module, "enforce_rbac_access") as enforce_rbac_access, + ): + unwrap(TagUpdateDeleteApi().delete)(TagUpdateDeleteApi(), request_context, "tag-1") + + tags_service.get_tag_type.assert_called_once_with(request_context, "tag-1") + enforce_rbac_access.assert_called_once() + + def test_delete_does_not_authorize_tag_outside_current_workspace( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.get_tag_type.return_value = None + tags_service.delete_tag.side_effect = TagNotFoundError() + owner = _account(TenantAccountRole.OWNER) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", True), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + patch.object(module, "enforce_rbac_access") as enforce_rbac_access, + ): + with pytest.raises(NotFound): + unwrap(TagUpdateDeleteApi().delete)(TagUpdateDeleteApi(), request_context, "tag-1") + + enforce_rbac_access.assert_not_called() -class TestTagBindingCollectionApi: - def test_create_success(self, app: Flask, admin_user, payload_patch): - api = TagBindingCollectionApi() - method = unwrap(api.post) +class TestTagBindings: + def test_create_passes_stable_binding_input( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + owner = _account(TenantAccountRole.OWNER) + payload = TagBindingPayload( + tag_ids=["tag-1", "tag-2"], + target_id="snippet-1", + type=TagType.SNIPPET, + ) - payload = { - "tag_ids": ["tag-1"], - "target_id": "target-1", - "type": "knowledge", - } + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + result, status = unwrap(TagBindingCollectionApi().post)(TagBindingCollectionApi(), payload, request_context) - with app.test_request_context("/", json=payload): - with ( - payload_patch(payload), - patch("controllers.console.tag.tags.TagService.save_tag_binding") as save_mock, - ): - result, status = method(api, TagBindingPayload.model_validate(payload), admin_user) + tags_service.create_bindings.assert_called_once_with( + request_context, + TagBindingInput(("tag-1", "tag-2"), "snippet-1", "snippet"), + ) + assert (result, status) == ({"result": "success"}, 200) - save_mock.assert_called_once() - assert status == 200 - assert result["result"] == "success" + def test_create_maps_missing_target_to_not_found( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.create_bindings.side_effect = TagBindingTargetNotFoundError("app") + owner = _account(TenantAccountRole.OWNER) + payload = TagBindingPayload(tag_ids=["tag-1"], target_id="missing", type=TagType.APP) - def test_create_snippet_binding_success(self, app: Flask, admin_user, payload_patch): - api = TagBindingCollectionApi() - method = unwrap(api.post) + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + with pytest.raises(NotFound, match="App not found") as exc_info: + unwrap(TagBindingCollectionApi().post)(TagBindingCollectionApi(), payload, request_context) - payload = { - "tag_ids": ["tag-1"], - "target_id": "snippet-1", - "type": "snippet", - } + assert exc_info.value.__cause__ is None + assert exc_info.value.__suppress_context__ is True - with app.test_request_context("/", json=payload): - with ( - payload_patch(payload), - patch("controllers.console.tag.tags.TagService.save_tag_binding") as save_mock, - ): - result, status = method(api, TagBindingPayload.model_validate(payload), admin_user) + def test_create_rejects_read_only_member(self, app: Flask, request_context: RequestContext) -> None: + readonly = _account(TenantAccountRole.NORMAL) + payload = TagBindingPayload(tag_ids=["tag-1"], target_id="app-1", type=TagType.APP) - save_mock.assert_called_once() - binding_payload = save_mock.call_args.args[0] - assert binding_payload.type == TagType.SNIPPET - assert binding_payload.target_id == "snippet-1" - assert status == 200 - assert result["result"] == "success" + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(readonly, "tenant-1")), + ): + with pytest.raises(Forbidden): + unwrap(TagBindingCollectionApi().post)(TagBindingCollectionApi(), payload, request_context) - def test_create_forbidden(self, app: Flask, readonly_user, payload_patch): - api = TagBindingCollectionApi() - method = unwrap(api.post) + def test_remove_passes_stable_binding_input( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + owner = _account(TenantAccountRole.OWNER) + payload = TagBindingRemovePayload( + tag_ids=["tag-1"], + target_id="app-1", + type=TagType.APP, + ) - with app.test_request_context("/", json={}): - with ( - payload_patch({}), - ): - with pytest.raises(Forbidden): - method( - api, - TagBindingPayload(tag_ids=["tag-1"], target_id="target-1", type=TagType.KNOWLEDGE), - readonly_user, - ) + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + result, status = unwrap(TagBindingRemoveApi().post)(TagBindingRemoveApi(), payload, request_context) + + tags_service.delete_bindings.assert_called_once_with( + request_context, + TagBindingInput(("tag-1",), "app-1", "app"), + ) + assert (result, status) == ({"result": "success"}, 200) + + def test_remove_maps_missing_target_to_not_found( + self, app: Flask, request_context: RequestContext, tags_service: MagicMock + ) -> None: + tags_service.delete_bindings.side_effect = TagBindingTargetNotFoundError("knowledge") + owner = _account(TenantAccountRole.OWNER) + payload = TagBindingRemovePayload(tag_ids=["tag-1"], target_id="missing", type=TagType.KNOWLEDGE) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(owner, "tenant-1")), + ): + with pytest.raises(NotFound, match="Dataset not found") as exc_info: + unwrap(TagBindingRemoveApi().post)(TagBindingRemoveApi(), payload, request_context) + + assert exc_info.value.__cause__ is None + assert exc_info.value.__suppress_context__ is True + + def test_remove_rejects_read_only_member(self, app: Flask, request_context: RequestContext) -> None: + readonly = _account(TenantAccountRole.NORMAL) + payload = TagBindingRemovePayload(tag_ids=["tag-1"], target_id="app-1", type=TagType.APP) + + with ( + app.test_request_context("/"), + patch.object(module.dify_config, "RBAC_ENABLED", False), + patch.object(module, "current_account_with_tenant", return_value=(readonly, "tenant-1")), + ): + with pytest.raises(Forbidden): + unwrap(TagBindingRemoveApi().post)(TagBindingRemoveApi(), payload, request_context) -class TestTagBindingRemoveApi: - def test_remove_success(self, app: Flask, admin_user, payload_patch): - api = TagBindingRemoveApi() - method = unwrap(api.post) - - payload = { - "tag_ids": ["tag-1", "tag-2"], - "target_id": "target-1", - "type": "knowledge", - } - - with app.test_request_context("/", json=payload): - with ( - payload_patch(payload), - patch("controllers.console.tag.tags.TagService.delete_tag_binding") as delete_mock, - ): - result, status = method(api, TagBindingRemovePayload.model_validate(payload), admin_user) - - delete_mock.assert_called_once() - delete_payload = delete_mock.call_args.args[0] - assert delete_payload.tag_ids == ["tag-1", "tag-2"] - assert status == 200 - assert result["result"] == "success" - - def test_remove_forbidden(self, app: Flask, readonly_user, payload_patch): - api = TagBindingRemoveApi() - method = unwrap(api.post) - - with app.test_request_context("/", json={}): - with ( - payload_patch({}), - ): - with pytest.raises(Forbidden): - method( - api, - TagBindingRemovePayload(tag_ids=["tag-1"], target_id="target-1", type=TagType.KNOWLEDGE), - readonly_user, - ) - - -class TestTagResponseModel: - def test_tag_response_normalizes_enum_type(self): +class TestTagResponseAndRoutes: + def test_tag_response_normalizes_enum_type(self) -> None: payload = module.TagResponse.model_validate( {"id": "tag-1", "name": "tag", "type": TagType.KNOWLEDGE, "binding_count": 1} ).model_dump(mode="json") @@ -473,32 +471,22 @@ class TestTagResponseModel: assert payload["type"] == "knowledge" assert payload["binding_count"] == "1" - -class TestTagBindingRouteMetadata: - def test_write_routes_are_not_deprecated(self): + def test_binding_routes_keep_contract(self) -> None: + assert TagBindingCollectionApi.post.__apidoc__["id"] == "create_tag_binding" + assert TagBindingRemoveApi.post.__apidoc__["id"] == "remove_tag_bindings" assert TagBindingCollectionApi.post.__apidoc__.get("deprecated") is not True assert TagBindingRemoveApi.post.__apidoc__.get("deprecated") is not True - def test_write_routes_have_stable_operation_ids(self): - assert TagBindingCollectionApi.post.__apidoc__["id"] == "create_tag_binding" - assert TagBindingRemoveApi.post.__apidoc__["id"] == "remove_tag_bindings" - - def test_write_routes_are_registered(self): route_map = { resource.__name__: urls for resource, urls, _route_doc, _kwargs in console_ns.resources - if resource.__name__ - in { - "TagBindingCollectionApi", - "TagBindingRemoveApi", - } + if resource.__name__ in {"TagBindingCollectionApi", "TagBindingRemoveApi"} + } + assert route_map == { + "TagBindingCollectionApi": ("/tag-bindings",), + "TagBindingRemoveApi": ("/tag-bindings/remove",), } - assert route_map["TagBindingCollectionApi"] == ("/tag-bindings",) - assert route_map["TagBindingRemoveApi"] == ("/tag-bindings/remove",) - - def test_legacy_write_routes_are_not_registered(self): urls = {url for _resource, resource_urls, _route_doc, _kwargs in console_ns.resources for url in resource_urls} - assert "/tag-bindings/create" not in urls assert "/tag-bindings/" not in urls diff --git a/api/tests/unit_tests/controllers/console/test_workspace_account.py b/api/tests/unit_tests/controllers/console/test_workspace_account.py index fce2993a815..5316b9efba4 100644 --- a/api/tests/unit_tests/controllers/console/test_workspace_account.py +++ b/api/tests/unit_tests/controllers/console/test_workspace_account.py @@ -8,6 +8,7 @@ import pytest from flask import Flask from sqlalchemy.orm import Session, scoped_session, sessionmaker +from controllers.console.error import EmailDomainSuspendedError from controllers.console.workspace.account import ( AccountDeleteUpdateFeedbackApi, ChangeEmailCheckApi, @@ -442,10 +443,29 @@ class TestChangeEmailValidity: class TestChangeEmailReset: + @patch( + "controllers.console.workspace.account.AccountService.get_account_freeze_type", + return_value="email_domain_suspended", + ) + def test_should_reject_suspended_email_domain(self, mock_get_freeze_type, app: Flask): + current_user = _build_account("old@example.com", "email-reset-account") + + with app.test_request_context( + "/account/change-email/reset", + method="POST", + json={"new_email": "new@suspended.example", "token": "token-123"}, + ): + api = ChangeEmailResetApi() + method = inspect.unwrap(api.post) + with pytest.raises(EmailDomainSuspendedError): + method(api, current_user) + + mock_get_freeze_type.assert_called_once_with("new@suspended.example") + @patch("controllers.console.workspace.account.AccountService.send_change_email_completed_notify_email") @patch("controllers.console.workspace.account.AccountService.revoke_change_email_token") @patch("controllers.console.workspace.account.AccountService.get_change_email_data") - @patch("controllers.console.workspace.account.AccountService.is_account_in_freeze") + @patch("controllers.console.workspace.account.AccountService.get_account_freeze_type") @pytest.mark.parametrize( "sqlite_session", [(Account, Tenant, TenantAccountJoin, AccountIntegrate)], @@ -507,7 +527,7 @@ class TestChangeEmailReset: @patch("controllers.console.workspace.account.AccountService.revoke_change_email_token") @patch("controllers.console.workspace.account.AccountService.get_change_email_data") @patch("controllers.console.workspace.account.AccountService.check_email_unique") - @patch("controllers.console.workspace.account.AccountService.is_account_in_freeze") + @patch("controllers.console.workspace.account.AccountService.get_account_freeze_type") def test_should_reject_reset_when_token_phase_is_not_new_verified( self, mock_is_freeze: MagicMock, @@ -550,7 +570,7 @@ class TestChangeEmailReset: @patch("controllers.console.workspace.account.AccountService.revoke_change_email_token") @patch("controllers.console.workspace.account.AccountService.get_change_email_data") @patch("controllers.console.workspace.account.AccountService.check_email_unique") - @patch("controllers.console.workspace.account.AccountService.is_account_in_freeze") + @patch("controllers.console.workspace.account.AccountService.get_account_freeze_type") def test_should_reject_reset_when_token_email_differs_from_payload_new_email( self, mock_is_freeze: MagicMock, @@ -593,7 +613,7 @@ class TestChangeEmailReset: @patch("controllers.console.workspace.account.AccountService.revoke_change_email_token") @patch("controllers.console.workspace.account.AccountService.get_change_email_data") @patch("controllers.console.workspace.account.AccountService.check_email_unique") - @patch("controllers.console.workspace.account.AccountService.is_account_in_freeze") + @patch("controllers.console.workspace.account.AccountService.get_account_freeze_type") def test_should_reject_reset_when_token_account_id_does_not_match_current_user( self, mock_is_freeze: MagicMock, @@ -736,7 +756,7 @@ class TestAccountDeletionFeedback: class TestCheckEmailUnique: - @patch("controllers.console.workspace.account.AccountService.is_account_in_freeze") + @patch("controllers.console.workspace.account.AccountService.get_account_freeze_type") @pytest.mark.parametrize( "sqlite_session", [(Account, Tenant, TenantAccountJoin)], diff --git a/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py b/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py index 321c957e79a..a46ce6b106d 100644 --- a/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py +++ b/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py @@ -5,9 +5,10 @@ import pytest from controllers.common.wraps import RBACPermission, RBACResourceScope from controllers.console.datasets.data_source import DataSourceApi +from controllers.console.datasets.rag_pipeline.datasource_auth import DatasourceAuth from controllers.console.workspace.model_providers import ModelProviderCredentialApi from controllers.console.workspace.models import ModelProviderModelCredentialApi -from controllers.console.workspace.tool_providers import ToolBuiltinProviderAddApi +from controllers.console.workspace.tool_providers import ToolBuiltinProviderAddApi, ToolOAuthCustomClient @pytest.mark.parametrize( @@ -50,3 +51,33 @@ def test_model_provider_credential_get_requires_admin_and_rbac( assert rbac_config["resource_type"] == RBACResourceScope.WORKSPACE assert rbac_config["scene"] == RBACPermission.CREDENTIAL_MANAGE assert rbac_config["resource_required"] is False + + +def test_tool_oauth_custom_client_get_requires_admin_and_rbac() -> None: + """GET endpoint that returns custom OAuth client params must enforce + the same admin + RBAC gates as its sibling POST and DELETE methods.""" + method = ToolOAuthCustomClient.get + + legacy_wrapper = unwrap(method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname) + assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname + + rbac_wrapper = unwrap(method, stop=lambda wrapper: "rbac_permission_required" in wrapper.__code__.co_qualname) + rbac_config = getclosurevars(rbac_wrapper).nonlocals + assert rbac_config["resource_type"] == RBACResourceScope.WORKSPACE + assert rbac_config["scene"] == RBACPermission.CREDENTIAL_MANAGE + assert rbac_config["resource_required"] is False + + +def test_datasource_auth_get_requires_edit_and_rbac() -> None: + """GET endpoint that lists datasource credentials must enforce + the same edit + RBAC gates as its sibling POST method.""" + method = DatasourceAuth.get + + edit_wrapper = unwrap(method, stop=lambda wrapper: "edit_permission_required" in wrapper.__code__.co_qualname) + assert "edit_permission_required" in edit_wrapper.__code__.co_qualname + + rbac_wrapper = unwrap(method, stop=lambda wrapper: "rbac_permission_required" in wrapper.__code__.co_qualname) + rbac_config = getclosurevars(rbac_wrapper).nonlocals + assert rbac_config["resource_type"] == RBACResourceScope.DATASET + assert rbac_config["scene"] == RBACPermission.CREDENTIAL_MANAGE + assert rbac_config["resource_required"] is False diff --git a/api/tests/unit_tests/controllers/console/workspace/test_accounts.py b/api/tests/unit_tests/controllers/console/workspace/test_accounts.py index d2857182858..f0480ea019f 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_accounts.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_accounts.py @@ -1,10 +1,12 @@ import inspect from datetime import UTC, datetime +from types import SimpleNamespace from unittest.mock import MagicMock, PropertyMock, patch from uuid import NAMESPACE_URL, uuid5 import pytest from flask import Flask +from jsonschema import Draft202012Validator from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound @@ -13,7 +15,7 @@ from controllers.console.auth.error import ( EmailAlreadyInUseError, EmailCodeError, ) -from controllers.console.error import AccountInFreezeError +from controllers.console.error import AccountInFreezeError, EmailDomainSuspendedError from controllers.console.workspace.account import ( AccountAvatarApi, AccountAvatarQuery, @@ -26,6 +28,7 @@ from controllers.console.workspace.account import ( AccountNameApi, AccountPasswordApi, AccountProfileApi, + AccountProfilePatchPayload, AccountTimezoneApi, ChangeEmailCheckApi, ChangeEmailResetApi, @@ -38,10 +41,12 @@ from controllers.console.workspace.error import ( ) from enums import DeploymentEdition from extensions.storage.storage_type import StorageType +from machinery.context import RequestContext from models import Account, AccountIntegrate, InvitationCode, Tenant, TenantAccountJoin from models.account import AccountStatus, InvitationCodeStatus, TenantAccountRole from models.enums import CreatorUserRole from models.model import UploadFile +from services.entities.account_entities import AccountProfileChanges from services.errors.account import CurrentPasswordIncorrectError as ServicePwdError @@ -165,28 +170,146 @@ class TestAccountProfileApi: class TestAccountUpdateApis: @pytest.mark.parametrize( - ("api_cls", "payload"), + ("api_cls", "payload", "expected_changes"), [ - (AccountNameApi, {"name": "test"}), - (AccountAvatarApi, {"avatar": "img.png"}), - (AccountInterfaceLanguageApi, {"interface_language": "en-US"}), - (AccountInterfaceThemeApi, {"interface_theme": "dark"}), - (AccountTimezoneApi, {"timezone": "UTC"}), + (AccountNameApi, {"name": "test"}, AccountProfileChanges(name="test")), + (AccountAvatarApi, {"avatar": "img.png"}, AccountProfileChanges(avatar="img.png")), + ( + AccountInterfaceLanguageApi, + {"interface_language": "en-US"}, + AccountProfileChanges(interface_language="en-US"), + ), + ( + AccountInterfaceThemeApi, + {"interface_theme": "dark"}, + AccountProfileChanges(interface_theme="dark"), + ), + (AccountTimezoneApi, {"timezone": "UTC"}, AccountProfileChanges(timezone="UTC")), ], ) - def test_update_success(self, app: Flask, api_cls, payload): + def test_deprecated_update_routes_delegate_to_profile_service( + self, app: Flask, api_cls, payload, expected_changes: AccountProfileChanges + ): api = api_cls() method = inspect.unwrap(api.post) - user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id=user.id, + active_workspace_id=None, + ) + profile = MagicMock() + profile.update.return_value = user with ( app.test_request_context("/", json=payload), - patch("controllers.console.workspace.account.AccountService.update_account", return_value=user), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(profile=profile)), + ), ): - result = method(api, user) + result = method(api, request_context) assert result["id"] == user.id + profile.update.assert_called_once_with(request_context, expected_changes) + + def test_deprecated_update_routes_are_marked_deprecated(self): + for api_cls in ( + AccountNameApi, + AccountAvatarApi, + AccountInterfaceLanguageApi, + AccountInterfaceThemeApi, + AccountTimezoneApi, + ): + assert api_cls.post.__apidoc__["deprecated"] is True + + +class TestAccountProfilePatchApi: + def test_json_schema_matches_runtime_patch_rules(self): + schema = AccountProfilePatchPayload.model_json_schema() + validator = Draft202012Validator(schema) + + assert schema["type"] == "object" + assert schema["additionalProperties"] is False + assert "required" not in schema + assert set(schema["properties"]) == { + "name", + "avatar", + "interface_language", + "interface_theme", + "timezone", + } + validator.validate({}) + validator.validate({"name": "Jane"}) + validator.validate({"name": "Jane", "interface_language": "en-US", "timezone": "UTC"}) + for payload in ( + {"name": None}, + {"unexpected": "value"}, + {"name": "Jane", "unexpected": "value"}, + ): + assert list(validator.iter_errors(payload)) + + def test_updates_multiple_profile_fields(self, app: Flask): + api = AccountProfileApi() + method = inspect.unwrap(api.patch) + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id=user.id, + active_workspace_id="workspace-1", + ) + profile = MagicMock() + profile.update.return_value = user + payload = {"name": "Jane", "interface_language": "en-US", "timezone": "UTC"} + args = AccountProfilePatchPayload.model_validate(payload) + + with ( + app.test_request_context("/account/profile", method="PATCH", json=payload), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(profile=profile)), + ), + ): + result = method(api, args, request_context) + + assert result["id"] == user.id + profile.update.assert_called_once_with( + request_context, + AccountProfileChanges(name="Jane", interface_language="en-US", timezone="UTC"), + ) + + def test_empty_patch_is_a_noop(self, app: Flask): + api = AccountProfileApi() + method = inspect.unwrap(api.patch) + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id=user.id, + active_workspace_id="workspace-1", + ) + profile = MagicMock() + profile.update.return_value = user + args = AccountProfilePatchPayload.model_validate({}) + + with ( + app.test_request_context("/account/profile", method="PATCH", json={}), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(profile=profile)), + ), + ): + result = method(api, args, request_context) + + assert result["id"] == user.id + profile.update.assert_called_once_with(request_context, AccountProfileChanges()) + + @pytest.mark.parametrize("payload", [{"name": None}, {"unexpected": "value"}]) + def test_rejects_null_or_unknown_changes(self, payload: dict[str, object]): + with pytest.raises(ValueError): + AccountProfilePatchPayload.model_validate(payload) class TestAccountAvatarApiGet: @@ -480,7 +603,7 @@ class TestChangeEmailApis: new_callable=PropertyMock, return_value=payload, ), - patch("controllers.console.workspace.account.AccountService.is_account_in_freeze", return_value=False), + patch("controllers.console.workspace.account.AccountService.get_account_freeze_type", return_value=None), patch("controllers.console.workspace.account.AccountService.check_email_unique", return_value=False), ): with pytest.raises(EmailAlreadyInUseError): @@ -502,7 +625,7 @@ class TestCheckEmailUniqueApi: new_callable=PropertyMock, return_value=payload, ), - patch("controllers.console.workspace.account.AccountService.is_account_in_freeze", return_value=False), + patch("controllers.console.workspace.account.AccountService.get_account_freeze_type", return_value=None), patch("controllers.console.workspace.account.AccountService.check_email_unique", return_value=True), ): result = method(api) @@ -523,7 +646,32 @@ class TestCheckEmailUniqueApi: new_callable=PropertyMock, return_value=payload, ), - patch("controllers.console.workspace.account.AccountService.is_account_in_freeze", return_value=True), + patch( + "controllers.console.workspace.account.AccountService.get_account_freeze_type", + return_value="freeze", + ), ): with pytest.raises(AccountInFreezeError): method(api) + + def test_email_domain_is_suspended(self, app: Flask): + api = CheckEmailUnique() + method = inspect.unwrap(api.post) + + payload = {"email": "user@suspended.example"} + + with ( + app.test_request_context("/", json=payload), + patch.object( + type(console_ns), + "payload", + new_callable=PropertyMock, + return_value=payload, + ), + patch( + "controllers.console.workspace.account.AccountService.get_account_freeze_type", + return_value="email_domain_suspended", + ), + ): + with pytest.raises(EmailDomainSuspendedError): + method(api) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py b/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py index a34349710ea..99f9ff1af9a 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py @@ -1,5 +1,6 @@ import inspect from datetime import UTC, datetime +from types import FunctionType from unittest.mock import patch import pytest @@ -23,6 +24,7 @@ from controllers.console.workspace.endpoint import ( EndpointUpdatePayload, LegacyEndpointUpdatePayload, ) +from controllers.console.wraps import RBACPermission, RBACResourceScope from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.plugin.entities.endpoint import EndpointEntityWithInstance, EndpointProviderDeclaration from core.plugin.impl.exc import PluginPermissionDeniedError @@ -56,6 +58,22 @@ def _endpoint_entity() -> EndpointEntityWithInstance: ) +@pytest.mark.parametrize("method", [EndpointListApi.get, EndpointListForSinglePluginApi.get]) +def test_endpoint_lists_require_management_permission(method: FunctionType) -> None: + legacy_wrapper = inspect.unwrap( + method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname + ) + assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname + + rbac_wrapper = inspect.unwrap( + method, stop=lambda wrapper: "rbac_permission_required" in wrapper.__code__.co_qualname + ) + rbac_config = inspect.getclosurevars(rbac_wrapper).nonlocals + assert rbac_config["resource_type"] == RBACResourceScope.WORKSPACE + assert rbac_config["scene"] == RBACPermission.PLUGIN_MODEL_CONFIG + assert rbac_config["resource_required"] is False + + class TestEndpointCollectionApi: def test_create_success(self, app: Flask): api = EndpointCollectionApi() diff --git a/api/tests/unit_tests/controllers/console/workspace/test_snippets.py b/api/tests/unit_tests/controllers/console/workspace/test_snippets.py index 211d9229a2a..e89aa6a6b57 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_snippets.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_snippets.py @@ -354,14 +354,40 @@ def test_export_snippet_returns_yaml_attachment(app: Flask, monkeypatch: pytest. api = snippets_module.CustomizedSnippetExportApi() handler = unwrap(api.get) - with app.test_request_context("/workspaces/current/customized-snippets/snippet-1/export?include_secret=true"): + with app.test_request_context( + "/workspaces/current/customized-snippets/snippet-1/export?include_secret=true&workflow_id=workflow-1" + ): response = handler(api, "tenant-1", snippet_id="snippet-1") assert response.status_code == 200 assert response.get_data(as_text=True) == "version: 0.1.0\nkind: snippet\n" assert response.headers["Content-Type"] == "application/x-yaml" assert "Snippet%20One.snippet" in response.headers["Content-Disposition"] - export_snippet_dsl.assert_called_once_with(snippet=snippet, include_secret=True) + export_snippet_dsl.assert_called_once_with(snippet=snippet, include_secret=True, workflow_id="workflow-1") + + +def test_export_snippet_raises_not_found_for_missing_workflow(app: Flask, monkeypatch: pytest.MonkeyPatch): + snippet = _snippet(name="Snippet One") + + monkeypatch.setattr(snippets_module.SnippetService, "get_snippet_by_id", Mock(return_value=snippet)) + monkeypatch.setattr( + snippets_module, + "SnippetDslService", + Mock( + return_value=SimpleNamespace( + export_snippet_dsl=Mock(side_effect=ValueError("Missing published workflow workflow-1")) + ) + ), + ) + monkeypatch.setattr(snippets_module, "Session", _SessionContext) + monkeypatch.setattr(snippets_module, "db", SimpleNamespace(engine=object())) + + api = snippets_module.CustomizedSnippetExportApi() + handler = unwrap(api.get) + + with app.test_request_context("/workspaces/current/customized-snippets/snippet-1/export?workflow_id=workflow-1"): + with pytest.raises(NotFound, match="Missing published workflow workflow-1"): + handler(api, "tenant-1", snippet_id="snippet-1") def test_import_snippet_returns_202_for_pending_confirmation(app: Flask, monkeypatch: pytest.MonkeyPatch): diff --git a/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py b/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py index e0a59fe2f7b..a7cf55008aa 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py +++ b/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py @@ -2,7 +2,6 @@ Unit tests for inner_api auth decorators """ -from unittest.mock import patch from uuid import NAMESPACE_URL, uuid5 import pytest @@ -11,7 +10,6 @@ from sqlalchemy import Engine, event from sqlalchemy.orm import Session from werkzeug.exceptions import HTTPException -from configs import dify_config from controllers.inner_api.wraps import ( billing_inner_api_only, enterprise_inner_api_only, @@ -24,6 +22,16 @@ from models.enums import EndUserType from models.model import EndUser +@pytest.fixture(autouse=True) +def _inner_api_config(config_overrides) -> None: + config_overrides( + INNER_API=True, + INNER_API_KEY="valid_key", + PLUGIN_DAEMON_KEY="plugin_key", + INNER_API_KEY_FOR_PLUGIN="valid_plugin_key", + ) + + def _stable_uuid(value: str) -> str: return str(uuid5(NAMESPACE_URL, value)) @@ -41,14 +49,12 @@ class TestBillingInnerApiOnly: # Act with app.test_request_context(headers={"X-Inner-Api-Key": "valid_key"}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - result = protected_view() + result = protected_view() # Assert assert result == "success" - def test_should_return_404_when_inner_api_disabled(self, app: Flask): + def test_should_return_404_when_inner_api_disabled(self, app: Flask, config_overrides): """Test that 404 is returned when INNER_API is disabled""" # Arrange @@ -57,11 +63,11 @@ class TestBillingInnerApiOnly: return "success" # Act & Assert + config_overrides(INNER_API=False) with app.test_request_context(): - with patch.object(dify_config, "INNER_API", False): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 404 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 404 def test_should_return_401_when_api_key_missing(self, app: Flask): """Test that 401 is returned when X-Inner-Api-Key header is missing""" @@ -73,11 +79,9 @@ class TestBillingInnerApiOnly: # Act & Assert with app.test_request_context(headers={}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 401 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 401 def test_should_return_401_when_api_key_invalid(self, app: Flask): """Test that 401 is returned when X-Inner-Api-Key header is invalid""" @@ -89,11 +93,9 @@ class TestBillingInnerApiOnly: # Act & Assert with app.test_request_context(headers={"X-Inner-Api-Key": "invalid_key"}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 401 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 401 class TestEnterpriseInnerApiOnly: @@ -109,14 +111,12 @@ class TestEnterpriseInnerApiOnly: # Act with app.test_request_context(headers={"X-Inner-Api-Key": "valid_key"}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - result = protected_view() + result = protected_view() # Assert assert result == "success" - def test_should_return_404_when_inner_api_disabled(self, app: Flask): + def test_should_return_404_when_inner_api_disabled(self, app: Flask, config_overrides): """Test that 404 is returned when INNER_API is disabled""" # Arrange @@ -125,11 +125,11 @@ class TestEnterpriseInnerApiOnly: return "success" # Act & Assert + config_overrides(INNER_API=False) with app.test_request_context(): - with patch.object(dify_config, "INNER_API", False): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 404 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 404 def test_should_return_401_when_api_key_missing(self, app: Flask): """Test that 401 is returned when X-Inner-Api-Key header is missing""" @@ -141,11 +141,9 @@ class TestEnterpriseInnerApiOnly: # Act & Assert with app.test_request_context(headers={}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 401 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 401 def test_should_return_401_when_api_key_invalid(self, app: Flask): """Test that 401 is returned when X-Inner-Api-Key header is invalid""" @@ -157,11 +155,9 @@ class TestEnterpriseInnerApiOnly: # Act & Assert with app.test_request_context(headers={"X-Inner-Api-Key": "invalid_key"}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 401 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 401 class TestInnerApiOnly: @@ -173,22 +169,20 @@ class TestInnerApiOnly: return "success" with app.test_request_context(headers={"X-Inner-Api-Key": "valid_key"}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - result = protected_view() + result = protected_view() assert result == "success" - def test_should_return_404_when_inner_api_disabled(self, app: Flask): + def test_should_return_404_when_inner_api_disabled(self, app: Flask, config_overrides): @inner_api_only def protected_view(): return "success" + config_overrides(INNER_API=False) with app.test_request_context(): - with patch.object(dify_config, "INNER_API", False): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 404 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 404 def test_should_return_401_when_api_key_missing(self, app: Flask): @inner_api_only @@ -196,11 +190,9 @@ class TestInnerApiOnly: return "success" with app.test_request_context(headers={}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 401 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 401 def test_should_return_401_when_api_key_invalid(self, app: Flask): @inner_api_only @@ -208,17 +200,15 @@ class TestInnerApiOnly: return "success" with app.test_request_context(headers={"X-Inner-Api-Key": "invalid_key"}): - with patch.object(dify_config, "INNER_API", True): - with patch.object(dify_config, "INNER_API_KEY", "valid_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 401 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 401 class TestEnterpriseInnerApiUserAuth: """Test enterprise_inner_api_user_auth decorator for HMAC-based user authentication""" - def test_should_pass_through_when_inner_api_disabled(self, app: Flask): + def test_should_pass_through_when_inner_api_disabled(self, app: Flask, config_overrides): """Test that request passes through when INNER_API is disabled""" # Arrange @@ -227,9 +217,9 @@ class TestEnterpriseInnerApiUserAuth: return kwargs.get("user", "no_user") # Act + config_overrides(INNER_API=False) with app.test_request_context(): - with patch.object(dify_config, "INNER_API", False): - result = protected_view() + result = protected_view() # Assert assert result == "no_user" @@ -244,8 +234,7 @@ class TestEnterpriseInnerApiUserAuth: # Act with app.test_request_context(headers={}): - with patch.object(dify_config, "INNER_API", True): - result = protected_view() + result = protected_view() # Assert assert result == "no_user" @@ -260,8 +249,7 @@ class TestEnterpriseInnerApiUserAuth: # Act with app.test_request_context(headers={"Authorization": "invalid_format"}): - with patch.object(dify_config, "INNER_API", True): - result = protected_view() + result = protected_view() # Assert assert result == "no_user" @@ -282,8 +270,7 @@ class TestEnterpriseInnerApiUserAuth: with app.test_request_context( headers={"Authorization": "Bearer user123:wrong_signature", "X-Inner-Api-Key": "valid_key"} ): - with patch.object(dify_config, "INNER_API", True): - result = protected_view() + result = protected_view() finally: event.remove(sqlite_engine, "before_cursor_execute", fail_on_query) @@ -322,8 +309,7 @@ class TestEnterpriseInnerApiUserAuth: with app.test_request_context( headers={"Authorization": f"Bearer {user_id}:{valid_signature}", "X-Inner-Api-Key": inner_api_key} ): - with patch.object(dify_config, "INNER_API", True): - result = protected_view() + result = protected_view() # Assert assert isinstance(result, EndUser) @@ -345,14 +331,12 @@ class TestPluginInnerApiOnly: # Act with app.test_request_context(headers={"X-Inner-Api-Key": "valid_plugin_key"}): - with patch.object(dify_config, "PLUGIN_DAEMON_KEY", "plugin_key"): - with patch.object(dify_config, "INNER_API_KEY_FOR_PLUGIN", "valid_plugin_key"): - result = protected_view() + result = protected_view() # Assert assert result == "success" - def test_should_return_404_when_plugin_daemon_key_not_set(self, app: Flask): + def test_should_return_404_when_plugin_daemon_key_not_set(self, app: Flask, config_overrides): """Test that 404 is returned when PLUGIN_DAEMON_KEY is not set""" # Arrange @@ -361,11 +345,11 @@ class TestPluginInnerApiOnly: return "success" # Act & Assert + config_overrides(PLUGIN_DAEMON_KEY="") with app.test_request_context(): - with patch.object(dify_config, "PLUGIN_DAEMON_KEY", ""): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 404 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 404 def test_should_return_404_when_api_key_invalid(self, app: Flask): """Test that 404 is returned when X-Inner-Api-Key header is invalid (note: returns 404, not 401)""" @@ -377,35 +361,31 @@ class TestPluginInnerApiOnly: # Act & Assert with app.test_request_context(headers={"X-Inner-Api-Key": "invalid_key"}): - with patch.object(dify_config, "PLUGIN_DAEMON_KEY", "plugin_key"): - with patch.object(dify_config, "INNER_API_KEY_FOR_PLUGIN", "valid_plugin_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() - assert exc_info.value.code == 404 + with pytest.raises(HTTPException) as exc_info: + protected_view() + assert exc_info.value.code == 404 class TestKnowledgeFSInnerApiOnly: """KnowledgeFS uses the same trusted transport key without exposing plugin semantics.""" - def test_should_allow_valid_shared_inner_key(self, app: Flask): + def test_should_allow_valid_shared_inner_key(self, app: Flask, config_overrides): @knowledge_fs_inner_api_only def protected_view(): return "success" + config_overrides(PLUGIN_DAEMON_KEY="plugin_key", INNER_API_KEY_FOR_PLUGIN="valid_plugin_key") with app.test_request_context(headers={"X-Inner-Api-Key": "valid_plugin_key"}): - with patch.object(dify_config, "PLUGIN_DAEMON_KEY", "plugin_key"): - with patch.object(dify_config, "INNER_API_KEY_FOR_PLUGIN", "valid_plugin_key"): - assert protected_view() == "success" + assert protected_view() == "success" - def test_should_hide_endpoint_for_invalid_key(self, app: Flask): + def test_should_hide_endpoint_for_invalid_key(self, app: Flask, config_overrides): @knowledge_fs_inner_api_only def protected_view(): return "success" + config_overrides(PLUGIN_DAEMON_KEY="plugin_key", INNER_API_KEY_FOR_PLUGIN="valid_plugin_key") with app.test_request_context(headers={"X-Inner-Api-Key": "invalid"}): - with patch.object(dify_config, "PLUGIN_DAEMON_KEY", "plugin_key"): - with patch.object(dify_config, "INNER_API_KEY_FOR_PLUGIN", "valid_plugin_key"): - with pytest.raises(HTTPException) as exc_info: - protected_view() + with pytest.raises(HTTPException) as exc_info: + protected_view() assert exc_info.value.code == 404 diff --git a/api/tests/unit_tests/controllers/openapi/test_app_dsl.py b/api/tests/unit_tests/controllers/openapi/test_app_dsl.py new file mode 100644 index 00000000000..25b0cd85917 --- /dev/null +++ b/api/tests/unit_tests/controllers/openapi/test_app_dsl.py @@ -0,0 +1,49 @@ +from inspect import unwrap +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from flask import Flask +from sqlalchemy.engine import Engine +from werkzeug.exceptions import Forbidden + +from controllers.openapi import app_dsl as app_dsl_module +from controllers.openapi._models import AppDslImportPayload +from controllers.openapi.app_dsl import AppDslImportApi, AppDslImportConfirmApi +from services.errors.account import NoPermissionError + + +@pytest.mark.parametrize( + ("api", "kwargs"), + [ + ( + AppDslImportApi(), + { + "workspace_id": "workspace-1", + "body": AppDslImportPayload(mode="yaml-content", yaml_content="app: {}"), + }, + ), + ( + AppDslImportConfirmApi(), + {"workspace_id": "workspace-1", "import_id": "import-1"}, + ), + ], +) +def test_permission_denial_maps_to_forbidden( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + api: AppDslImportApi | AppDslImportConfirmApi, + kwargs: dict[str, object], +) -> None: + service = Mock() + service.import_app.side_effect = NoPermissionError("denied") + service.confirm_import.side_effect = NoPermissionError("denied") + monkeypatch.setattr(app_dsl_module, "AppDslService", Mock(return_value=service)) + monkeypatch.setattr(app_dsl_module, "db", SimpleNamespace(engine=sqlite_engine)) + + with app.test_request_context("/openapi/v1/workspaces/workspace-1/apps/imports", method="POST"): + with pytest.raises(Forbidden, match="denied") as exc_info: + unwrap(api.post)(api, auth_data=SimpleNamespace(caller=Mock()), **kwargs) + + assert isinstance(exc_info.value.__cause__, NoPermissionError) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_app.py b/api/tests/unit_tests/controllers/service_api/app/test_app.py index a6ff026f6d5..38d7eb7e678 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_app.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_app.py @@ -312,6 +312,7 @@ def test_get_site_configuration_queries_authenticated_app( input_placeholder="Ask anything", custom_disclaimer=None, default_language="en-US", + prompt_public=False, show_workflow_steps=True, use_icon_as_answer_icon=False, ) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_audio.py b/api/tests/unit_tests/controllers/service_api/app/test_audio.py index 4fb56e6ac8b..091c129c874 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_audio.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_audio.py @@ -10,7 +10,6 @@ Tests coverage for: import io import uuid from inspect import unwrap -from types import SimpleNamespace from unittest.mock import Mock, patch import pytest @@ -34,6 +33,8 @@ from controllers.service_api.app.error import ( ) from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError +from models.enums import EndUserType +from models.model import App, AppMode, EndUser from services.app_ref_service import AppRef, MessageRef from services.audio_service import AudioService from services.errors.app_model_config import AppModelConfigBrokenError @@ -50,6 +51,31 @@ def _file_data(): return FileStorage(stream=io.BytesIO(b"audio"), filename="audio.wav", content_type="audio/wav") +def _app(*, app_id: str = "a1", tenant_id: str = "tenant-1") -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name="Audio app", + description="", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + + +def _end_user(*, end_user_id: str = "u1", external_user_id: str | None = None) -> EndUser: + return EndUser( + id=end_user_id, + tenant_id="tenant-1", + app_id="a1", + type=EndUserType.SERVICE_API, + external_user_id=external_user_id, + name="Audio user", + session_id=f"session-{end_user_id}", + ) + + # --------------------------------------------------------------------------- # Pydantic Model Tests # --------------------------------------------------------------------------- @@ -197,8 +223,8 @@ class TestAudioApi: monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "ok"}) api = AudioApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") - end_user = SimpleNamespace(id="u1") + app_model = _app() + end_user = _end_user() with app.test_request_context("/audio-to-text", method="POST", data={"file": _file_data()}): response = handler(api, app_model=app_model, end_user=end_user) @@ -224,8 +250,8 @@ class TestAudioApi: monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(exc)) api = AudioApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") - end_user = SimpleNamespace(id="u1") + app_model = _app() + end_user = _end_user() with app.test_request_context("/audio-to-text", method="POST", data={"file": _file_data()}): with pytest.raises(expected): @@ -237,8 +263,8 @@ class TestAudioApi: ) api = AudioApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") - end_user = SimpleNamespace(id="u1") + app_model = _app() + end_user = _end_user() with app.test_request_context("/audio-to-text", method="POST", data={"file": _file_data()}): with pytest.raises(InternalServerError): @@ -251,8 +277,8 @@ class TestTextApi: api = TextApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") - end_user = SimpleNamespace(id="end-user-1", external_user_id="ext") + app_model = _app() + end_user = _end_user(end_user_id="end-user-1", external_user_id="ext") with app.test_request_context( "/text-to-audio", @@ -274,8 +300,8 @@ class TestTextApi: api = TextApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1", tenant_id="tenant-1") - end_user = SimpleNamespace(id="end-user-1", external_user_id="ext") + app_model = _app() + end_user = _end_user(end_user_id="end-user-1", external_user_id="ext") with app.test_request_context( "/text-to-audio", @@ -294,8 +320,8 @@ class TestTextApi: api = TextApi() handler = unwrap(api.post) - app_model = SimpleNamespace(id="a1") - end_user = SimpleNamespace(id="end-user-1", external_user_id="ext") + app_model = _app() + end_user = _end_user(end_user_id="end-user-1", external_user_id="ext") with app.test_request_context("/text-to-audio", method="POST", json={"text": "hello"}): with pytest.raises(ProviderQuotaExceededError): diff --git a/api/tests/unit_tests/controllers/test_swagger.py b/api/tests/unit_tests/controllers/test_swagger.py index 20bdeacda3c..9ea27878c8e 100644 --- a/api/tests/unit_tests/controllers/test_swagger.py +++ b/api/tests/unit_tests/controllers/test_swagger.py @@ -6,6 +6,12 @@ from collections.abc import Iterator import pytest from flask import Flask + +@pytest.fixture(autouse=True) +def _swagger_config(config_overrides) -> None: + config_overrides(SWAGGER_UI_ENABLED=True) + + USER_PROPERTY_SCHEMA = { "description": ( "User identifier, unique within the application. This identifier scopes data access; resources created with " @@ -154,14 +160,11 @@ def test_uuid_path_format_is_derived_from_route_converter(): } -def test_openapi_json_endpoints_render(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_openapi_json_endpoints_render(): from controllers.console import bp as console_bp from controllers.service_api import bp as service_api_bp from controllers.web import bp as web_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -190,12 +193,9 @@ def test_openapi_json_endpoints_render(monkeypatch: pytest.MonkeyPatch): assert app.config["RESTX_INCLUDE_ALL_MODELS"] is True -def test_service_document_file_routes_document_multipart_form_data(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_document_file_routes_document_multipart_form_data(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -236,12 +236,9 @@ def test_service_document_file_routes_document_multipart_form_data(monkeypatch: assert update_operation["requestBody"]["required"] is False -def test_service_openapi_merges_public_api_reference_descriptions(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_merges_public_api_reference_descriptions(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -261,12 +258,9 @@ def test_service_openapi_merges_public_api_reference_descriptions(monkeypatch: p assert _parameters_by_name(rename_operation)["c_id"]["description"] == "Conversation ID." -def test_service_document_list_documents_query_params_render(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_document_list_documents_query_params_render(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -280,12 +274,9 @@ def test_service_document_list_documents_query_params_render(monkeypatch: pytest assert params[name]["in"] == "query" -def test_service_openapi_documents_decorator_user_contracts(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_documents_decorator_user_contracts(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -330,12 +321,9 @@ def test_service_openapi_documents_decorator_user_contracts(monkeypatch: pytest. assert events_params["user"]["required"] is True -def test_service_openapi_documents_app_multipart_contracts(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_documents_app_multipart_contracts(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -365,12 +353,9 @@ def test_service_openapi_documents_app_multipart_contracts(monkeypatch: pytest.M assert pipeline_schema["required"] == ["file"] -def test_service_openapi_documents_non_json_response_media_types(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_documents_non_json_response_media_types(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -409,12 +394,9 @@ def test_service_openapi_documents_non_json_response_media_types(monkeypatch: py } -def test_service_openapi_documents_uuid_params_and_deprecated_routes(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_documents_uuid_params_and_deprecated_routes(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -441,12 +423,9 @@ def test_service_openapi_documents_uuid_params_and_deprecated_routes(monkeypatch assert paths["/datasets/{dataset_id}/documents/{document_id}/update_by_text"]["post"]["deprecated"] is True -def test_service_openapi_documents_path_action_enums(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_documents_path_action_enums(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -465,12 +444,9 @@ def test_service_openapi_documents_path_action_enums(monkeypatch: pytest.MonkeyP assert metadata_params["action"]["schema"]["enum"] == ["enable", "disable"] -def test_service_openapi_documents_conditional_payload_schemas(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_documents_conditional_payload_schemas(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -497,12 +473,9 @@ def test_service_openapi_documents_conditional_payload_schemas(monkeypatch: pyte assert without_text_branch["properties"]["text"]["type"] == "null" -def test_service_openapi_does_not_encode_docs_coverage_boundaries(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_does_not_encode_docs_coverage_boundaries(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -525,12 +498,9 @@ def test_service_openapi_does_not_encode_docs_coverage_boundaries(monkeypatch: p assert paths["/datasets/{dataset_id}/documents/{document_id}/update-by-file"]["post"]["deprecated"] is True -def test_service_openapi_documents_auth_and_compatibility_payloads(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_service_openapi_documents_auth_and_compatibility_payloads(): from controllers.service_api import bp as service_api_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -556,12 +526,9 @@ def test_service_openapi_documents_auth_and_compatibility_payloads(monkeypatch: assert tag_ids_schema["required"] == ["tag_ids", "target_id"] -def test_console_account_avatar_query_param_renders_as_query(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_console_account_avatar_query_param_renders_as_query(): from controllers.console import bp as console_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True @@ -576,11 +543,39 @@ def test_console_account_avatar_query_param_renders_as_query(monkeypatch: pytest assert params["avatar"]["required"] is True -def test_console_agent_debug_conversation_refresh_has_no_body(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_console_account_profile_patch_and_deprecated_aliases(): from controllers.console import bp as console_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) + app = Flask(__name__) + app.config["TESTING"] = True + app.config["RESTX_INCLUDE_ALL_MODELS"] = True + app.register_blueprint(console_bp) + + payload = app.test_client().get("/console/api/openapi.json").get_json() + paths = payload["paths"] + + profile_patch = paths["/account/profile"]["patch"] + assert profile_patch.get("deprecated") is not True + profile_patch_schema = _json_body_schema(payload, profile_patch) + assert profile_patch_schema["type"] == "object" + assert profile_patch_schema["additionalProperties"] is False + assert "required" not in profile_patch_schema + assert profile_patch_schema["properties"]["name"]["type"] == "string" + + for path in ( + "/account/name", + "/account/avatar", + "/account/interface-language", + "/account/interface-theme", + "/account/timezone", + ): + assert paths[path]["post"]["deprecated"] is True + + assert paths["/account/avatar"]["get"].get("deprecated") is not True + + +def test_console_agent_debug_conversation_refresh_has_no_body(): + from controllers.console import bp as console_bp app = Flask(__name__) app.config["TESTING"] = True @@ -593,12 +588,9 @@ def test_console_agent_debug_conversation_refresh_has_no_body(monkeypatch: pytes assert "AgentDebugConversationRefreshPayload" not in payload["components"]["schemas"] -def test_console_member_invite_documents_bad_request_response(monkeypatch: pytest.MonkeyPatch): - from configs import dify_config +def test_console_member_invite_documents_bad_request_response(): from controllers.console import bp as console_bp - monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True) - app = Flask(__name__) app.config["TESTING"] = True app.config["RESTX_INCLUDE_ALL_MODELS"] = True diff --git a/api/tests/unit_tests/controllers/web/test_app.py b/api/tests/unit_tests/controllers/web/test_app.py index 3c5daf9224c..005c23e7f25 100644 --- a/api/tests/unit_tests/controllers/web/test_app.py +++ b/api/tests/unit_tests/controllers/web/test_app.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest from flask import Flask +from werkzeug.exceptions import Unauthorized from controllers.common.errors import InvalidArgumentError from controllers.web.app import AppAccessMode, AppMeta, AppParameterApi, AppWebAuthPermission @@ -14,6 +15,7 @@ from controllers.web.error import ( AgentNotPublishedError, AppUnavailableError, WebAppAccessServiceUnavailableError, + WebAppAuthRequiredError, WebAppNotFoundError, ) from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict @@ -194,14 +196,151 @@ class TestAppAccessMode: # AppWebAuthPermission # --------------------------------------------------------------------------- class TestAppWebAuthPermission: - @patch("controllers.web.app.WebAppAuthService.is_app_require_permission_check", return_value=False) - def test_returns_true_when_no_permission_check_required(self, mock_check: MagicMock, app: Flask) -> None: - with app.test_request_context("/webapp/permission?appId=app-1", headers={"X-App-Code": "code1"}): + @patch("controllers.web.app.application_services") + def test_returns_true_without_reading_passport_when_no_permission_check_required( + self, application_services: MagicMock, app: Flask + ) -> None: + webapp_access = MagicMock() + webapp_access.requires_permission_check.return_value = False + application_services.return_value = SimpleNamespace(webapp_access=webapp_access) + + with ( + app.test_request_context("/webapp/permission?appId=app-1", headers={"X-App-Code": "code1"}), + patch("controllers.web.app.extract_webapp_passport") as extract_passport, + ): result = AppWebAuthPermission().get() assert result == {"result": True} + webapp_access.requires_permission_check.assert_called_once_with("app-1") + webapp_access.is_user_allowed.assert_not_called() + extract_passport.assert_not_called() - def test_raises_when_missing_app_id(self, app: Flask) -> None: - with app.test_request_context("/webapp/permission", headers={"X-App-Code": "code1"}): + @pytest.mark.parametrize( + ("decoded", "expected_user_id", "allowed"), + [ + pytest.param({"user_id": "user-1"}, "user-1", True, id="identified-user"), + pytest.param({}, "visitor", False, id="visitor-fallback"), + ], + ) + @patch("controllers.web.app.application_services") + def test_checks_private_app_permission( + self, + application_services: MagicMock, + decoded: dict[str, str], + expected_user_id: str, + allowed: bool, + app: Flask, + ) -> None: + webapp_access = MagicMock() + webapp_access.requires_permission_check.return_value = True + webapp_access.is_user_allowed.return_value = allowed + application_services.return_value = SimpleNamespace(webapp_access=webapp_access) + + with ( + app.test_request_context("/webapp/permission?appId=app-1", headers={"X-App-Code": "code1"}), + patch("controllers.web.app.extract_webapp_passport", return_value="passport") as extract_passport, + patch("controllers.web.app.PassportService") as passport_service, + ): + passport_service.return_value.verify.return_value = decoded + result = AppWebAuthPermission().get() + + assert result == {"result": allowed} + webapp_access.requires_permission_check.assert_called_once_with("app-1") + extract_passport.assert_called_once() + passport_service.return_value.verify.assert_called_once_with("passport") + webapp_access.is_user_allowed.assert_called_once_with(user_id=expected_user_id, app_id="app-1") + + @pytest.mark.parametrize("failing_method", ["requires_permission_check", "is_user_allowed"]) + @patch("controllers.web.app.application_services") + def test_maps_access_dependency_failure_to_service_unavailable( + self, application_services: MagicMock, failing_method: str, app: Flask + ) -> None: + webapp_access = MagicMock() + webapp_access.requires_permission_check.return_value = True + if failing_method == "requires_permission_check": + webapp_access.requires_permission_check.side_effect = WebAppAccessUnavailableError() + else: + webapp_access.is_user_allowed.side_effect = WebAppAccessUnavailableError() + application_services.return_value = SimpleNamespace(webapp_access=webapp_access) + + passport_service = MagicMock() + passport_service.return_value.verify.return_value = {"user_id": "user-1"} + with ( + app.test_request_context("/webapp/permission?appId=app-1", headers={"X-App-Code": "code1"}), + patch("controllers.web.app.extract_webapp_passport", return_value="passport"), + patch("controllers.web.app.PassportService", passport_service), + pytest.raises(WebAppAccessServiceUnavailableError) as raised, + ): + AppWebAuthPermission().get() + + assert raised.value.data == { + "code": "web_app_access_unavailable", + "message": "Web app access service is unavailable.", + "status": 503, + } + + @patch("controllers.web.app.application_services") + def test_private_app_requires_passport(self, application_services: MagicMock, app: Flask) -> None: + webapp_access = MagicMock() + webapp_access.requires_permission_check.return_value = True + application_services.return_value = SimpleNamespace(webapp_access=webapp_access) + + with ( + app.test_request_context("/webapp/permission?appId=app-1", headers={"X-App-Code": "code1"}), + patch("controllers.web.app.extract_webapp_passport", return_value=None), + pytest.raises(WebAppAuthRequiredError) as raised, + ): + AppWebAuthPermission().get() + + assert raised.value.data == { + "code": "web_sso_auth_required", + "message": "Web app authentication required.", + "status": 401, + } + webapp_access.is_user_allowed.assert_not_called() + + @pytest.mark.parametrize( + "description", + ["Token has expired.", "Invalid token signature.", "Invalid token."], + ) + @patch("controllers.web.app.application_services") + def test_invalid_passport_is_normalized_to_web_app_auth_required( + self, application_services: MagicMock, description: str, app: Flask + ) -> None: + webapp_access = MagicMock() + webapp_access.requires_permission_check.return_value = True + application_services.return_value = SimpleNamespace(webapp_access=webapp_access) + invalid_passport = Unauthorized(description) + + with ( + app.test_request_context("/webapp/permission?appId=app-1", headers={"X-App-Code": "code1"}), + patch("controllers.web.app.extract_webapp_passport", return_value="passport"), + patch("controllers.web.app.PassportService") as passport_service, + ): + passport_service.return_value.verify.side_effect = invalid_passport + with pytest.raises(WebAppAuthRequiredError) as raised: + AppWebAuthPermission().get() + + assert raised.value.data == { + "code": "web_sso_auth_required", + "message": "Web app authentication required.", + "status": 401, + } + webapp_access.is_user_allowed.assert_not_called() + + @pytest.mark.parametrize( + ("path", "headers"), + [ + pytest.param("/webapp/permission", {"X-App-Code": "code1"}, id="missing-app-id"), + pytest.param("/webapp/permission?appId=app-1", {}, id="missing-app-code"), + ], + ) + @patch("controllers.web.app.application_services") + def test_raises_when_app_reference_is_missing( + self, application_services: MagicMock, path: str, headers: dict[str, str], app: Flask + ) -> None: + with app.test_request_context(path, headers=headers): with pytest.raises(ValueError, match="appId"): AppWebAuthPermission().get() + + application_services.assert_not_called() diff --git a/api/tests/unit_tests/controllers/web/test_human_input_form.py b/api/tests/unit_tests/controllers/web/test_human_input_form.py index 985c1e41ba9..5d9ba9e42e4 100644 --- a/api/tests/unit_tests/controllers/web/test_human_input_form.py +++ b/api/tests/unit_tests/controllers/web/test_human_input_form.py @@ -14,7 +14,6 @@ from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import Forbidden import controllers.web.human_input_form as human_input_module -import controllers.web.site as site_module from controllers.web.error import WebFormRateLimitExceededError from core.workflow.nodes.human_input.entities import ParagraphInputConfig, SelectInputConfig, StringListSource from core.workflow.nodes.human_input.enums import ValueSourceType @@ -139,7 +138,7 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask, dat monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock) monkeypatch.setattr( - site_module.FeatureService, + human_input_module.FeatureService, "get_features", lambda tenant_id, **_kwargs: FeatureModel(can_replace_logo=True, webapp_copyright_enabled=True), ) @@ -259,7 +258,7 @@ def test_get_form_uses_runtime_select_options(monkeypatch: pytest.MonkeyPatch, a def mock_get_features(tenant_id: str, exclude_vector_space: bool = False): return FeatureModel(can_replace_logo=True) - monkeypatch.setattr(site_module.FeatureService, "get_features", mock_get_features) + monkeypatch.setattr(human_input_module.FeatureService, "get_features", mock_get_features) with app.test_request_context("/api/form/human_input/token-1", method="GET"): response = HumanInputFormApi().get("token-1") @@ -360,7 +359,7 @@ def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: F monkeypatch.setattr(human_input_module, "HumanInputService", lambda engine: service_mock) monkeypatch.setattr( - site_module.FeatureService, + human_input_module.FeatureService, "get_features", lambda tenant_id, **_kwargs: FeatureModel(can_replace_logo=True, webapp_copyright_enabled=True), ) diff --git a/api/tests/unit_tests/controllers/web/test_site.py b/api/tests/unit_tests/controllers/web/test_site.py index 011d6b6a51e..8440796984a 100644 --- a/api/tests/unit_tests/controllers/web/test_site.py +++ b/api/tests/unit_tests/controllers/web/test_site.py @@ -1,106 +1,82 @@ +from types import SimpleNamespace from unittest.mock import MagicMock, patch -from configs import dify_config +import pytest +from werkzeug.exceptions import Forbidden + from controllers.web import site as site_module -from enums import DeploymentEdition -from extensions.storage.storage_type import StorageType -from models.model import AppMode, IconType, Site -from services.entities.feature_entities import FeatureModel +from services.app_definition_query_service import AppSiteConfiguration +from services.web_app_runtime_query_service import WebAppBootstrap, WebAppRuntimeUnavailableError -def test_app_site_api_returns_legacy_agent_compatible_mode() -> None: - app_model = MagicMock() - app_model.id = "app-id" - app_model.tenant_id = "tenant-id" - app_model.tenant = MagicMock(id="tenant-id", status="normal") - app_model.mode_compatible_with_agent_with_session.return_value = AppMode.AGENT_CHAT +def _bootstrap() -> WebAppBootstrap: + site = AppSiteConfiguration( + title="Test Site", + chat_color_theme="light", + chat_color_theme_inverted=False, + icon_type="image", + icon="file-1", + icon_background="#ffffff", + description="Description", + copyright="Copyright", + privacy_policy="Privacy", + input_placeholder="Ask anything", + custom_disclaimer="Disclaimer", + default_language="en-US", + prompt_public=True, + show_workflow_steps=True, + use_icon_as_answer_icon=False, + ) + return WebAppBootstrap( + app_id="app-id", + mode="agent-chat", + enable_site=True, + site={**site._asdict(), "icon_url": "https://files.example.com/icon.png"}, + plan="pro", + can_replace_logo=True, + custom_config={ + "remove_webapp_brand": True, + "replace_webapp_logo": "https://files.example.com/files/workspaces/tenant-id/webapp-logo", + }, + ) + + +def test_app_site_api_queries_the_admitted_app_runtime() -> None: + app_model = MagicMock(id="app-id") end_user = MagicMock(id="end-user-id") - site = Site() - response = MagicMock() - response.model_dump.return_value = {"mode": AppMode.AGENT_CHAT} + web_app_runtime = MagicMock() + web_app_runtime.get_bootstrap.return_value = _bootstrap() - with ( - patch.object(site_module, "db") as mock_db, - patch.object(site_module.FeatureService, "get_features", return_value=FeatureModel(can_replace_logo=False)), - patch.object(site_module, "_build_site_icon_url", return_value=None), - patch.object(site_module.WebAppSiteResponse, "from_app_site", return_value=response) as mock_from_app_site, + with patch.object( + site_module, + "application_services", + return_value=SimpleNamespace(web_app_runtime=web_app_runtime), ): - mock_db.session.scalar.return_value = site result = site_module.AppSiteApi().get(app_model, end_user) - assert result["mode"] == AppMode.AGENT_CHAT - app_model.mode_compatible_with_agent_with_session.assert_called_once_with(session=mock_db.session()) - mock_from_app_site.assert_called_once_with( - tenant=app_model.tenant, - app_model=app_model, - mode=AppMode.AGENT_CHAT, - site=site, - end_user_id=end_user.id, - features=FeatureModel(can_replace_logo=False), - can_replace_logo=False, - icon_url=None, - ) + assert result["app_id"] == "app-id" + assert result["mode"] == "agent-chat" + assert result["end_user_id"] == "end-user-id" + assert result["site"]["prompt_public"] is True + assert result["site"]["icon_url"] == "https://files.example.com/icon.png" + assert result["model_config"] is None + assert result["custom_config"] == { + "remove_webapp_brand": True, + "replace_webapp_logo": "https://files.example.com/files/workspaces/tenant-id/webapp-logo", + } + web_app_runtime.get_bootstrap.assert_called_once_with("app-id") -def test_build_site_icon_url_uses_s3_presigned_url() -> None: - site = Site( - icon_type=IconType.IMAGE, - icon="11111111-1111-4111-8111-111111111111", - ) +def test_app_site_api_maps_unavailable_runtime_to_forbidden() -> None: + web_app_runtime = MagicMock() + web_app_runtime.get_bootstrap.side_effect = WebAppRuntimeUnavailableError with ( - patch.object(dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.CLOUD), - patch.object(dify_config, "STORAGE_TYPE", StorageType.S3), - patch.object(site_module, "db") as mock_db, - patch.object(site_module, "FileService") as mock_file_service, - patch.object(site_module, "build_icon_url") as mock_build_icon_url, + patch.object( + site_module, + "application_services", + return_value=SimpleNamespace(web_app_runtime=web_app_runtime), + ), + pytest.raises(Forbidden), ): - mock_file_service.return_value.get_file_presigned_url.return_value = ( - "https://s3.example.com/icon.png?signature=test" - ) - - result = site_module._build_site_icon_url(site=site, tenant_id="tenant-id") - - assert result == "https://s3.example.com/icon.png?signature=test" - mock_file_service.assert_called_once_with(mock_db.engine) - mock_file_service.return_value.get_file_presigned_url.assert_called_once_with( - file_id="11111111-1111-4111-8111-111111111111", - tenant_id="tenant-id", - ) - mock_build_icon_url.assert_not_called() - - -def test_build_site_icon_url_keeps_preview_url_for_self_hosted_s3() -> None: - site = Site( - icon_type=IconType.IMAGE, - icon="11111111-1111-4111-8111-111111111111", - ) - - with ( - patch.object(dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY), - patch.object(dify_config, "STORAGE_TYPE", StorageType.S3), - patch.object(site_module, "FileService") as mock_file_service, - patch.object(site_module, "build_icon_url", return_value="https://api.example.com/files/icon/file-preview"), - ): - result = site_module._build_site_icon_url(site=site, tenant_id="tenant-id") - - assert result == "https://api.example.com/files/icon/file-preview" - mock_file_service.assert_not_called() - - -def test_build_site_icon_url_keeps_preview_url_for_non_s3_storage() -> None: - site = Site( - icon_type=IconType.IMAGE, - icon="11111111-1111-4111-8111-111111111111", - ) - - with ( - patch.object(dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.CLOUD), - patch.object(dify_config, "STORAGE_TYPE", StorageType.LOCAL), - patch.object(site_module, "FileService") as mock_file_service, - patch.object(site_module, "build_icon_url", return_value="https://api.example.com/files/icon/file-preview"), - ): - result = site_module._build_site_icon_url(site=site, tenant_id="tenant-id") - - assert result == "https://api.example.com/files/icon/file-preview" - mock_file_service.assert_not_called() + site_module.AppSiteApi().get(MagicMock(id="app-id"), MagicMock(id="end-user-id")) diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py index 2e4c996798e..776d14aab9b 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py @@ -706,6 +706,52 @@ class TestEasyUiBasedGenerateTaskPipeline: assert isinstance(responses[-1].err, ValueError) assert pipeline._task_state.llm_result.message.content == "annotated" + def test_process_stream_response_error_event_adds_trace_task(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) + application_generate_entity.extras = {"trace_session_id": "session-1"} + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + _set_queue_events(pipeline, [_queue_message(QueueErrorEvent(error=ValueError("boom")))]) + _set_method(pipeline, "handle_error", lambda **kwargs: ValueError("boom")) + _set_method(pipeline, "error_to_stream_response", lambda err: ErrorStreamResponse(task_id="task", err=err)) + + trace_manager_double = _TraceManagerDouble() + trace_manager = cast(TraceQueueManager, trace_manager_double) + + class _Session: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def commit(self): + return None + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) + + assert len(responses) == 1 + assert isinstance(responses[0], ErrorStreamResponse) + trace_manager_double.add_trace_task.assert_called_once() + trace_task = trace_manager_double.add_trace_task.call_args.args[0] + assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE + assert trace_task.conversation_id == "conv" + assert trace_task.message_id == "msg" + assert trace_task.kwargs["trace_session_id"] == "session-1" + def test_agent_thought_to_stream_response_returns_payload(self, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) message = _make_message() diff --git a/api/tests/unit_tests/core/plugin/test_plugin_runtime.py b/api/tests/unit_tests/core/plugin/test_plugin_runtime.py index 704b82adc00..e3dcf2996b4 100644 --- a/api/tests/unit_tests/core/plugin/test_plugin_runtime.py +++ b/api/tests/unit_tests/core/plugin/test_plugin_runtime.py @@ -77,13 +77,12 @@ class TestPluginRuntimeExecution: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-api-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-api-key", + ) def test_request_preparation(self, plugin_client, mock_config): """Test that requests are properly prepared with correct headers and URL.""" @@ -182,13 +181,12 @@ class TestPluginRuntimeSandboxIsolation: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "secure-api-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="secure-api-key", + ) def test_api_key_authentication(self, plugin_client, mock_config): """Test that all requests include API key for authentication.""" @@ -272,13 +270,13 @@ class TestPluginRuntimeResourceLimits: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration with timeout.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - patch("core.plugin.impl.base.plugin_daemon_request_timeout", httpx.Timeout(30.0)), - ): + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + with patch("core.plugin.impl.base.plugin_daemon_request_timeout", httpx.Timeout(30.0)): yield def test_timeout_configuration_applied(self, plugin_client, mock_config): @@ -346,13 +344,12 @@ class TestPluginRuntimeErrorHandling: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_plugin_invoke_rate_limit_error(self, plugin_client, mock_config): """Test handling of rate limit errors during plugin invocation.""" @@ -605,13 +602,12 @@ class TestPluginRuntimeCommunication: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_request_response_communication(self, plugin_client, mock_config): """Test basic request/response communication pattern.""" @@ -811,13 +807,12 @@ class TestPluginToolManagerIntegration: return PluginToolManager() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_tool_invocation_success(self, tool_manager, mock_config): """Test successful tool invocation.""" @@ -938,13 +933,12 @@ class TestPluginInstallerIntegration: return PluginInstaller() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_list_plugins_success(self, installer, mock_config): """Test successful plugin listing.""" @@ -1012,13 +1006,12 @@ class TestPluginRuntimeEdgeCases: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_malformed_json_response(self, plugin_client, mock_config): """Test handling of malformed JSON responses.""" @@ -1174,13 +1167,12 @@ class TestPluginRuntimeAdvancedScenarios: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_multiple_sequential_requests(self, plugin_client, mock_config): """Test multiple sequential requests to the same endpoint.""" @@ -1360,13 +1352,12 @@ class TestPluginRuntimeSecurityAndValidation: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "secure-key-123"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="secure-key-123", + ) def test_api_key_header_always_present(self, plugin_client, mock_config): """Test that API key header is always included in requests.""" @@ -1480,13 +1471,12 @@ class TestPluginRuntimePerformanceScenarios: return BasePluginClient() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_high_volume_streaming(self, plugin_client, mock_config): """Test streaming with high volume of chunks.""" @@ -1598,13 +1588,12 @@ class TestPluginToolManagerAdvanced: return PluginToolManager() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_tool_invocation_with_complex_parameters(self, tool_manager, mock_config): """Test tool invocation with complex parameter structures.""" @@ -1750,13 +1739,12 @@ class TestPluginInstallerAdvanced: return PluginInstaller() @pytest.fixture - def mock_config(self): + def mock_config(self, config_overrides): """Mock plugin daemon configuration.""" - with ( - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_URL", "http://127.0.0.1:5002"), - patch("core.plugin.impl.base.dify_config.PLUGIN_DAEMON_KEY", "test-key"), - ): - yield + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) def test_upload_plugin_package_success(self, installer, mock_config): """Test successful plugin package upload.""" diff --git a/api/tests/unit_tests/core/tools/test_signature.py b/api/tests/unit_tests/core/tools/test_signature.py index 4985fc16368..7dc87014935 100644 --- a/api/tests/unit_tests/core/tools/test_signature.py +++ b/api/tests/unit_tests/core/tools/test_signature.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from typing import Literal from urllib.parse import parse_qs, urlparse @@ -18,6 +19,16 @@ from core.tools.signature import ( ) +@pytest.fixture(autouse=True) +def _signature_config(config_overrides: Callable[..., None]) -> None: + config_overrides( + SECRET_KEY="unit-secret", + FILES_URL="https://files.example.com", + INTERNAL_FILES_URL="https://internal.example.com", + FILES_ACCESS_TIMEOUT=120, + ) + + def test_bind_file_uri_uses_selected_base_and_preserves_remote_url() -> None: uri = "/files/tools/tool-file-id.png?sign=1" @@ -31,7 +42,6 @@ def test_bind_file_uri_uses_selected_base_and_preserves_remote_url() -> None: def test_sign_tool_file_uri_has_no_origin(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x08" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") uri = sign_tool_file_uri("tool-file-id", ".png") parsed = urlparse(uri) @@ -45,10 +55,6 @@ def test_sign_tool_file_uri_has_no_origin(monkeypatch: pytest.MonkeyPatch) -> No def test_sign_tool_file_and_verify_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x01" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_URL", "https://files.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.INTERNAL_FILES_URL", "https://internal.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 120) url = sign_tool_file("tool-file-id", ".png", for_external=False) parsed = urlparse(url) @@ -66,10 +72,6 @@ def test_sign_tool_file_and_verify_roundtrip(monkeypatch: pytest.MonkeyPatch) -> def test_sign_tool_file_for_external_uses_files_url(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x04" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_URL", "https://files.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.INTERNAL_FILES_URL", "https://internal.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 120) url = sign_tool_file("tool-file-id", ".png", for_external=True) parsed = urlparse(url) @@ -79,13 +81,12 @@ def test_sign_tool_file_for_external_uses_files_url(monkeypatch: pytest.MonkeyPa assert parsed.path == "/files/tools/tool-file-id.png" -def test_verify_tool_file_signature_rejects_invalid_sign(monkeypatch: pytest.MonkeyPatch) -> None: +def test_verify_tool_file_signature_rejects_invalid_sign( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(INTERNAL_FILES_URL="", FILES_ACCESS_TIMEOUT=10) monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x02" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_URL", "https://files.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.INTERNAL_FILES_URL", "") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 10) url = sign_tool_file("tool-file-id", ".txt") parsed = urlparse(url) @@ -97,13 +98,12 @@ def test_verify_tool_file_signature_rejects_invalid_sign(monkeypatch: pytest.Mon assert verify_tool_file_signature("tool-file-id", timestamp, nonce, "bad-signature") is False -def test_verify_tool_file_signature_rejects_expired_signature(monkeypatch: pytest.MonkeyPatch) -> None: +def test_verify_tool_file_signature_rejects_expired_signature( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(INTERNAL_FILES_URL="", FILES_ACCESS_TIMEOUT=10) monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x02" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_URL", "https://files.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.INTERNAL_FILES_URL", "") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 10) url = sign_tool_file("tool-file-id", ".txt") parsed = urlparse(url) @@ -119,9 +119,6 @@ def test_verify_tool_file_signature_rejects_expired_signature(monkeypatch: pytes def test_sign_upload_file_preview_url_uses_files_url(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x03" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_URL", "https://files.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.INTERNAL_FILES_URL", "https://internal.example.com") url = sign_upload_file_preview_url("upload-id", ".png") parsed = urlparse(url) @@ -137,9 +134,6 @@ def test_sign_upload_file_preview_url_uses_files_url(monkeypatch: pytest.MonkeyP def test_sign_upload_file_preview_url_ignores_internal_files_url(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x05" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_URL", "https://files.example.com") - monkeypatch.setattr("core.tools.signature.dify_config.INTERNAL_FILES_URL", "https://internal.example.com") url = sign_upload_file_preview_url("upload-id", ".png") parsed = urlparse(url) @@ -152,11 +146,12 @@ def test_sign_upload_file_preview_url_ignores_internal_files_url(monkeypatch: py assert query["sign"][0] -def test_get_signed_file_uri_for_plugin_and_verify_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None: +def test_get_signed_file_uri_for_plugin_and_verify_roundtrip( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(FILES_ACCESS_TIMEOUT=60) monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x06" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 60) uri = get_signed_file_uri_for_plugin( filename="report.pdf", @@ -198,13 +193,13 @@ def test_get_signed_file_uri_for_plugin_and_verify_roundtrip(monkeypatch: pytest ) def test_plugin_upload_signature_binds_max_size_without_legacy_payload_ambiguity( monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], user_from: Literal["account", "end-user"] | None, forged_nonce_suffix: str, ) -> None: + config_overrides(FILES_ACCESS_TIMEOUT=60) monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x0a" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 60) uri = get_signed_file_uri_for_plugin( filename="report.pdf", @@ -233,11 +228,12 @@ def test_plugin_upload_signature_binds_max_size_without_legacy_payload_ambiguity assert verify_plugin_file_signature(**forged) is False -def test_plugin_upload_signature_binds_account_user_from(monkeypatch: pytest.MonkeyPatch) -> None: +def test_plugin_upload_signature_binds_account_user_from( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(FILES_ACCESS_TIMEOUT=60) monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x09" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 60) uri = get_signed_file_uri_for_plugin( filename="report.pdf", @@ -263,11 +259,12 @@ def test_plugin_upload_signature_binds_account_user_from(monkeypatch: pytest.Mon assert verify_plugin_file_signature(**signed) is False -def test_verify_plugin_file_signature_rejects_invalid_signatures(monkeypatch: pytest.MonkeyPatch) -> None: +def test_verify_plugin_file_signature_rejects_invalid_signatures( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(FILES_ACCESS_TIMEOUT=30) monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x07" * 16) - monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") - monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 30) uri = get_signed_file_uri_for_plugin( filename="report.pdf", diff --git a/api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py b/api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py index f3903e2e438..ccb0cdff7a0 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py +++ b/api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py @@ -10,8 +10,6 @@ Test coverage: - Disabled mode behavior """ -from unittest.mock import patch - import pytest from opentelemetry.trace import StatusCode @@ -21,10 +19,14 @@ from graphon.enums import BuiltinNodeTypes from graphon.graph_events import GraphRunAbortedEvent +@pytest.fixture(autouse=True) +def _otel_config(config_overrides) -> None: + config_overrides(ENABLE_OTEL=True) + + class TestObservabilityLayerInitialization: """Test ObservabilityLayer initialization logic.""" - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_initialization_when_otel_enabled(self, tracer_provider_with_memory_exporter): """Test that layer initializes correctly when OTel is enabled.""" @@ -34,10 +36,10 @@ class TestObservabilityLayerInitialization: assert BuiltinNodeTypes.TOOL in layer._parsers assert layer._default_parser is not None - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", False) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_true") - def test_initialization_when_instrument_flag_enabled(self, tracer_provider_with_memory_exporter): + def test_initialization_when_instrument_flag_enabled(self, tracer_provider_with_memory_exporter, config_overrides): """Test that layer enables when instrument flag is enabled.""" + config_overrides(ENABLE_OTEL=False) layer = ObservabilityLayer() assert not layer._is_disabled assert layer._tracer is not None @@ -48,7 +50,6 @@ class TestObservabilityLayerInitialization: class TestObservabilityLayerNodeSpanLifecycle: """Test node span creation and lifecycle management.""" - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_node_span_created_and_ended( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node @@ -65,7 +66,6 @@ class TestObservabilityLayerNodeSpanLifecycle: assert spans[0].name == mock_llm_node.title assert spans[0].status.status_code == StatusCode.OK - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_node_error_recorded_in_span( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node @@ -84,7 +84,6 @@ class TestObservabilityLayerNodeSpanLifecycle: assert len(spans[0].events) > 0 assert any("exception" in event.name.lower() for event in spans[0].events) - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_node_end_without_start_handled_gracefully( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node @@ -102,7 +101,6 @@ class TestObservabilityLayerNodeSpanLifecycle: class TestObservabilityLayerParserIntegration: """Test parser integration for different node types.""" - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_default_parser_used_for_regular_node( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node @@ -121,7 +119,6 @@ class TestObservabilityLayerParserIntegration: assert attrs["node.execution_id"] == mock_start_node.execution_id assert attrs["node.type"] == mock_start_node.node_type - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_tool_parser_used_for_tool_node( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_tool_node @@ -140,7 +137,6 @@ class TestObservabilityLayerParserIntegration: assert attrs["gen_ai.tool.name"] == mock_tool_node.title assert attrs["gen_ai.tool.type"] == mock_tool_node._node_data.provider_type.value - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_llm_parser_used_for_llm_node( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node, mock_result_event @@ -178,7 +174,6 @@ class TestObservabilityLayerParserIntegration: assert attrs["gen_ai.completion"] == "test completion" assert attrs["gen_ai.response.finish_reason"] == "stop" - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_retrieval_parser_used_for_retrieval_node( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_retrieval_node, mock_result_event @@ -206,7 +201,6 @@ class TestObservabilityLayerParserIntegration: assert attrs["retrieval.query"] == "test query" assert "retrieval.document" in attrs - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_result_event_extracts_inputs_and_outputs( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, mock_result_event @@ -237,7 +231,6 @@ class TestObservabilityLayerParserIntegration: class TestObservabilityLayerGraphLifecycle: """Test graph lifecycle management.""" - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_on_graph_start_clears_contexts(self, tracer_provider_with_memory_exporter, mock_llm_node): """Test that on_graph_start clears node contexts.""" @@ -250,7 +243,6 @@ class TestObservabilityLayerGraphLifecycle: layer.on_graph_start() assert len(layer._node_contexts) == 0 - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_on_graph_end_with_no_unfinished_spans( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node @@ -266,7 +258,6 @@ class TestObservabilityLayerGraphLifecycle: spans = memory_span_exporter.get_finished_spans() assert len(spans) == 1 - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_on_graph_end_with_unfinished_spans_logs_warning( self, tracer_provider_with_memory_exporter, mock_llm_node, caplog @@ -283,7 +274,6 @@ class TestObservabilityLayerGraphLifecycle: assert len(layer._node_contexts) == 0 assert "node spans were not properly ended" in caplog.text - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", True) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") def test_graph_aborted_event_records_reason_on_current_span( self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node @@ -308,10 +298,10 @@ class TestObservabilityLayerGraphLifecycle: class TestObservabilityLayerDisabledMode: """Test behavior when layer is disabled.""" - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", False) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_disabled_mode_skips_node_start(self, memory_span_exporter, mock_start_node): + def test_disabled_mode_skips_node_start(self, memory_span_exporter, mock_start_node, config_overrides): """Test that disabled layer doesn't create spans on node start.""" + config_overrides(ENABLE_OTEL=False) layer = ObservabilityLayer() assert layer._is_disabled @@ -322,10 +312,10 @@ class TestObservabilityLayerDisabledMode: spans = memory_span_exporter.get_finished_spans() assert len(spans) == 0 - @patch("core.app.workflow.layers.observability.dify_config.ENABLE_OTEL", False) @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_disabled_mode_skips_node_end(self, memory_span_exporter, mock_llm_node): + def test_disabled_mode_skips_node_end(self, memory_span_exporter, mock_llm_node, config_overrides): """Test that disabled layer doesn't process node end.""" + config_overrides(ENABLE_OTEL=False) layer = ObservabilityLayer() assert layer._is_disabled diff --git a/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py b/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py index d8d3c1ecc2e..ee094732b65 100644 --- a/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py @@ -441,6 +441,48 @@ def test_extract_text_from_excel_numeric_type_column(mock_excel_file): assert expected_manual == result +@pytest.mark.parametrize( + ("extension", "mime_type", "route_label"), + [ + (".odt", "text/plain", "extension"), + (None, "application/vnd.oasis.opendocument.text", "mime_type"), + ], +) +def test_extract_text_from_file_routes_odt_inputs_to_graphon_odt_extractor( + document_extractor_node, + extension, + mime_type, + route_label, +): + file = Mock(spec=File) + file.extension = extension + file.mime_type = mime_type + + def fake_partition(file_content, *, suffix, unstructured_api_config, load_local_partition, render_element): + assert file_content == b"odt content" + assert suffix == ".odt" + assert unstructured_api_config == document_extractor_node._unstructured_api_config + assert load_local_partition.__name__ == "_load_partition_odt" + assert render_element is not None + return f"extracted through {route_label}" + + with ( + patch( + "graphon.nodes.document_extractor.node._download_file_content", + return_value=b"odt content", + ) as mock_download, + patch("graphon.nodes.document_extractor.node._partition_unstructured_file", side_effect=fake_partition), + ): + text = _extract_text_from_file( + document_extractor_node.http_client, + file, + unstructured_api_config=document_extractor_node._unstructured_api_config, + ) + + assert text == f"extracted through {route_label}" + mock_download.assert_called_once_with(document_extractor_node.http_client, file) + + @pytest.mark.parametrize( ("extension", "mime_type"), [ diff --git a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py index 4d760748839..3e9d23d14e4 100644 --- a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py @@ -1,8 +1,14 @@ +import datetime import json +import time +from collections.abc import Generator from unittest.mock import MagicMock, patch +import pytest + from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( LogstoreAPIWorkflowNodeExecutionRepository, + _dict_to_workflow_node_execution_model, ) from models.workflow import WorkflowNodeExecutionModel @@ -39,3 +45,45 @@ def test_get_execution_by_id_keeps_process_data_from_highest_failed_log_version( assert execution is not None assert execution.status.value == "failed" assert execution.process_data_dict == {"workflow_agent_binding_id": "binding-1"} + + +_CREATED_AT = datetime.datetime(2026, 8, 18, 2, 0, 0, tzinfo=datetime.UTC) +_FINISHED_AT = _CREATED_AT + datetime.timedelta(seconds=30) + + +@pytest.fixture +def non_utc_host_timezone(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: + """Run the host clock in UTC+05:30 so local-time conversions become observable.""" + monkeypatch.setenv("TZ", "Asia/Kolkata") + time.tzset() + yield + monkeypatch.undo() + time.tzset() + + +@pytest.mark.parametrize( + ("case", "payload"), + [ + ("both epoch", {"created_at": _CREATED_AT.timestamp(), "finished_at": _FINISHED_AT.timestamp()}), + ("aware iso and epoch", {"created_at": _CREATED_AT.isoformat(), "finished_at": _FINISHED_AT.timestamp()}), + ( + "naive iso and epoch", + {"created_at": _CREATED_AT.replace(tzinfo=None).isoformat(), "finished_at": _FINISHED_AT.timestamp()}, + ), + ("both datetime", {"created_at": _CREATED_AT, "finished_at": _FINISHED_AT}), + ], +) +@pytest.mark.usefixtures("non_utc_host_timezone") +def test_dict_to_node_execution_normalizes_timestamps_to_naive_utc(case: str, payload: dict[str, object]) -> None: + model = _dict_to_workflow_node_execution_model({"id": "execution-1", **payload}) + + assert model.created_at == _CREATED_AT.replace(tzinfo=None), case + assert model.finished_at == _FINISHED_AT.replace(tzinfo=None), case + + +@pytest.mark.usefixtures("non_utc_host_timezone") +def test_dict_to_node_execution_defaults_missing_created_at_to_naive_utc_now() -> None: + model = _dict_to_workflow_node_execution_model({"id": "execution-1"}) + + assert model.created_at.tzinfo is None + assert abs((model.created_at - datetime.datetime.now(tz=datetime.UTC).replace(tzinfo=None)).total_seconds()) < 60 diff --git a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_run_repository.py b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_run_repository.py new file mode 100644 index 00000000000..f8fd9d41eab --- /dev/null +++ b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_run_repository.py @@ -0,0 +1,54 @@ +import datetime +import time +from collections.abc import Generator + +import pytest + +from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run + +_START = datetime.datetime(2026, 8, 18, 2, 0, 0, tzinfo=datetime.UTC) +_FINISH = _START + datetime.timedelta(seconds=30) +_EXPECTED_CREATED_AT = _START.replace(tzinfo=None) +_EXPECTED_FINISHED_AT = _FINISH.replace(tzinfo=None) + +_BASE: dict[str, object] = {"id": "run-1", "tenant_id": "tenant-1", "app_id": "app-1", "workflow_id": "workflow-1"} + + +@pytest.fixture +def non_utc_host_timezone(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: + """Run the host clock in UTC+05:30 so local-time conversions become observable.""" + monkeypatch.setenv("TZ", "Asia/Kolkata") + time.tzset() + yield + monkeypatch.undo() + time.tzset() + + +@pytest.mark.parametrize( + ("case", "payload"), + [ + ("both epoch", {"started_at": _START.timestamp(), "finished_at": _FINISH.timestamp()}), + ("aware iso and epoch", {"started_at": _START.isoformat(), "finished_at": _FINISH.timestamp()}), + ( + "naive iso and epoch", + {"started_at": _START.replace(tzinfo=None).isoformat(), "finished_at": _FINISH.timestamp()}, + ), + ("both datetime", {"started_at": _START, "finished_at": _FINISH}), + ], +) +@pytest.mark.usefixtures("non_utc_host_timezone") +def test_dict_to_workflow_run_normalizes_timestamps_to_naive_utc(case: str, payload: dict[str, object]) -> None: + model = _dict_to_workflow_run({**_BASE, **payload}) + + assert model.created_at == _EXPECTED_CREATED_AT, case + assert model.finished_at == _EXPECTED_FINISHED_AT, case + assert model.elapsed_time == 30.0, case + + +@pytest.mark.usefixtures("non_utc_host_timezone") +def test_dict_to_workflow_run_defaults_missing_started_at_to_naive_utc_now() -> None: + model = _dict_to_workflow_run({**_BASE, "finished_at": _FINISH.timestamp()}) + + assert model.created_at.tzinfo is None + # A naive local-time default would sit 5h30m ahead of UTC and drive elapsed_time negative. + assert abs((model.created_at - datetime.datetime.now(tz=datetime.UTC).replace(tzinfo=None)).total_seconds()) < 60 diff --git a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_execution_repository.py b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_execution_repository.py new file mode 100644 index 00000000000..0e598e657fe --- /dev/null +++ b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_execution_repository.py @@ -0,0 +1,39 @@ +from collections.abc import Callable +from types import SimpleNamespace +from typing import cast +from unittest.mock import patch + +from sqlalchemy.orm import Session, sessionmaker + +from extensions.logstore.repositories.logstore_workflow_execution_repository import ( + LogstoreWorkflowExecutionRepository, +) +from models.account import Account +from models.enums import WorkflowRunTriggeredFrom + + +def test_repository_uses_typed_logstore_migration_flags( + config_overrides: Callable[..., None], + sqlite_session_factory: sessionmaker[Session], +) -> None: + config_overrides( + LOGSTORE_DUAL_WRITE_ENABLED=True, + LOGSTORE_ENABLE_PUT_GRAPH_FIELD=False, + ) + with ( + patch("extensions.logstore.repositories.logstore_workflow_execution_repository.AliyunLogStore"), + patch( + "extensions.logstore.repositories.logstore_workflow_execution_repository." + "SQLAlchemyWorkflowExecutionRepository" + ), + ): + repository = LogstoreWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=cast(Account, SimpleNamespace(id="account-1")), + app_id="app-1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + assert repository._enable_dual_write is True + assert repository._enable_put_graph_field is False diff --git a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py index fbde472a719..63d87ba6966 100644 --- a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py @@ -1,6 +1,6 @@ +from collections.abc import Callable from unittest.mock import MagicMock, patch -import pytest from sqlalchemy.orm import Session, sessionmaker from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( @@ -17,9 +17,9 @@ def _make_account() -> Account: def test_save_synchronously_writes_sql_when_dual_write_is_disabled( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + config_overrides: Callable[..., None], sqlite_session_factory: sessionmaker[Session] ) -> None: - monkeypatch.delenv("LOGSTORE_DUAL_WRITE_ENABLED", raising=False) + config_overrides(LOGSTORE_DUAL_WRITE_ENABLED=False) with ( patch("extensions.logstore.repositories.logstore_workflow_node_execution_repository.AliyunLogStore"), patch( diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index 5ddaeb99412..d727f47221f 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -3,18 +3,22 @@ import json from types import SimpleNamespace from unittest.mock import MagicMock, patch +from uuid import uuid4 import httpx import pytest from flask import Flask from pydantic import ValidationError +from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from enums import DeploymentEdition, WebAppAccessMode from extensions import ext_application_services from extensions.ext_redis import RedisClientWrapper -from models.model import DifySetup +from models.model import AccountTrialAppRecord, DifySetup from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository +from repositories.account_repository import SQLAlchemyAccountRepository +from services import recommended_app_catalog_gateway from services.account_activation_adapters import ( BillingAccountActivationEligibility, BillingWorkspaceMembershipCache, @@ -25,6 +29,7 @@ from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthS from services.enterprise.enterprise_service import WebAppSettings from services.errors.enterprise import EnterpriseAPIError, EnterpriseAPINotFoundError from services.init_validation_service import InvalidInitializationPasswordError +from services.tag_application_service import TagApplicationService from services.webapp_access_query_service import WebAppAccessUnavailableError @@ -151,6 +156,34 @@ def test_build_application_services_does_not_construct_schema_manager( schema_manager.assert_not_called() +def test_build_application_services_wires_tag_boundary( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + + assert isinstance(services.tags, TagApplicationService) + + +def test_build_application_services_wires_account_profile_repository( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + + accounts = services.accounts.profile._accounts + assert isinstance(accounts, SQLAlchemyAccountRepository) + assert accounts._session_factory is sqlite_session_factory + + @pytest.mark.parametrize( ("deployment_edition", "billing_enabled"), [ @@ -195,6 +228,31 @@ def test_build_application_services_wires_data_source_api_key_auth( assert isinstance(services.data_source_api_key_auth, DataSourceApiKeyAuthService) +def test_build_application_services_wires_trial_app_usage( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + app_id = str(uuid4()) + account_id = str(uuid4()) + + services.trial_app_usage.record(app_id=app_id, account_id=account_id) + + with sqlite_session_factory() as session: + record = session.scalar( + select(AccountTrialAppRecord).where( + AccountTrialAppRecord.app_id == app_id, + AccountTrialAppRecord.account_id == account_id, + ) + ) + assert record is not None + assert record.count == 1 + + def test_build_application_services_adapts_enterprise_webapp_access_mode( sqlite_session_factory: sessionmaker[Session], ) -> None: @@ -300,3 +358,86 @@ def test_build_application_services_does_not_hide_unknown_enterprise_errors( services.webapp_access.get_access_mode(app_id="app-1", app_code=None) assert raised.value is failure + + +def test_build_application_services_wires_webapp_permission( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with ( + patch( + "extensions.ext_application_services.FeatureService.is_webapp_auth_enabled", return_value=True + ) as enabled, + patch( + "extensions.ext_application_services.EnterpriseService.WebAppAuth.get_app_access_mode_by_id", + return_value=SimpleNamespace(access_mode="private"), + ) as get_access_mode, + patch( + "extensions.ext_application_services.EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp", + return_value=False, + ) as is_user_allowed, + ): + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + requires_permission = services.webapp_access.requires_permission_check("app-1") + allowed = services.webapp_access.is_user_allowed(user_id="user-1", app_id="app-1") + + assert requires_permission is True + assert allowed is False + enabled.assert_called_once_with() + get_access_mode.assert_called_once_with("app-1") + is_user_allowed.assert_called_once_with("user-1", "app-1") + + +def test_webapp_permission_adapter_maps_connection_failure() -> None: + failure = httpx.ConnectError("connection failed") + with ( + patch( + "extensions.ext_application_services.EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp", + side_effect=failure, + ), + pytest.raises(WebAppAccessUnavailableError) as raised, + ): + ext_application_services._is_user_allowed_to_access_webapp("user-1", "app-1") + + assert raised.value.__cause__ is failure + + +def test_build_application_services_wires_dynamic_recommended_catalog( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(ext_application_services.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "builtin") + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + + builtin_payload = json.dumps( + { + "recommended_apps": { + "en-US": { + "recommended_apps": [{"app": None, "app_id": "app-1", "categories": []}], + "categories": [], + } + } + } + ) + with patch.object(recommended_app_catalog_gateway.Path, "read_text", return_value=builtin_payload): + result = services.recommended_app_queries.list_recommended( + requested_language="en-US", + interface_language=None, + ) + assert result.recommended_apps + + monkeypatch.setattr(ext_application_services.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "invalid") + with pytest.raises(ValueError, match="invalid fetch recommended apps mode: invalid"): + services.recommended_app_queries.list_recommended( + requested_language="en-US", + interface_language=None, + ) diff --git a/api/tests/unit_tests/libs/test_token.py b/api/tests/unit_tests/libs/test_token.py index f129a1d86ce..09e213651d4 100644 --- a/api/tests/unit_tests/libs/test_token.py +++ b/api/tests/unit_tests/libs/test_token.py @@ -41,26 +41,32 @@ def test_extract_access_token(): assert extract_webapp_access_token(request) == expected_webapp -def test_real_cookie_name_uses_host_prefix_without_domain(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "https://console.example.com", raising=False) - monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "https://api.example.com", raising=False) - monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", "", raising=False) +def test_real_cookie_name_uses_host_prefix_without_domain(config_overrides): + config_overrides( + CONSOLE_WEB_URL="https://console.example.com", + CONSOLE_API_URL="https://api.example.com", + COOKIE_DOMAIN="", + ) assert token._real_cookie_name("csrf_token") == "__Host-csrf_token" -def test_real_cookie_name_without_host_prefix_when_domain_present(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "https://console.example.com", raising=False) - monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "https://api.example.com", raising=False) - monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", ".example.com", raising=False) +def test_real_cookie_name_without_host_prefix_when_domain_present(config_overrides): + config_overrides( + CONSOLE_WEB_URL="https://console.example.com", + CONSOLE_API_URL="https://api.example.com", + COOKIE_DOMAIN=".example.com", + ) assert token._real_cookie_name("csrf_token") == "csrf_token" -def test_set_csrf_cookie_includes_domain_when_configured(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "https://console.example.com", raising=False) - monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "https://api.example.com", raising=False) - monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", ".example.com", raising=False) +def test_set_csrf_cookie_includes_domain_when_configured(config_overrides): + config_overrides( + CONSOLE_WEB_URL="https://console.example.com", + CONSOLE_API_URL="https://api.example.com", + COOKIE_DOMAIN=".example.com", + ) response = Response() request = MagicMock() @@ -94,12 +100,14 @@ def test_non_whitelisted_path_requires_csrf(): token.check_csrf_token(request, "account-1") -def test_admin_api_key_header_bypasses_csrf_when_console_cookie_is_present(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(token.dify_config, "ADMIN_API_KEY_ENABLE", True) - monkeypatch.setattr(token.dify_config, "ADMIN_API_KEY", "admin-key") - monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "http://console.example.com") - monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "http://api.example.com") - monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", "") +def test_admin_api_key_header_bypasses_csrf_when_console_cookie_is_present(config_overrides): + config_overrides( + ADMIN_API_KEY_ENABLE=True, + ADMIN_API_KEY="admin-key", + CONSOLE_WEB_URL="http://console.example.com", + CONSOLE_API_URL="http://api.example.com", + COOKIE_DOMAIN="", + ) request = cast( Request, MockRequest( diff --git a/api/tests/unit_tests/migrations/test_clean_legacy_agent_soul_files.py b/api/tests/unit_tests/migrations/test_clean_legacy_agent_soul_files.py new file mode 100644 index 00000000000..1f14a44c16a --- /dev/null +++ b/api/tests/unit_tests/migrations/test_clean_legacy_agent_soul_files.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import importlib.util +import json +from io import StringIO +from pathlib import Path +from types import ModuleType + +import pytest +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from pydantic import ValidationError + +from models.agent_config_entities import AgentSoulConfig + +_MIGRATION_PATH = ( + Path(__file__).resolve().parents[3] + / "migrations/versions/2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py" +) + + +def _load_migration_module() -> ModuleType: + spec = importlib.util.spec_from_file_location("clean_legacy_agent_soul_files", _MIGRATION_PATH) + if spec is None or spec.loader is None: + raise RuntimeError("failed to load migration module") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _run_upgrade(module: ModuleType, engine: sa.Engine) -> None: + with engine.begin() as connection: + operations = Operations(MigrationContext.configure(connection)) + original_op = module.__dict__["op"] + module.__dict__["op"] = operations + try: + module.__dict__["upgrade"]() + finally: + module.__dict__["op"] = original_op + + +def test_upgrade_makes_legacy_build_draft_valid_for_first_message() -> None: + engine = sa.create_engine("sqlite:///:memory:") + metadata = sa.MetaData() + for table_name in ("agent_config_snapshots", "agent_config_drafts"): + sa.Table( + table_name, + metadata, + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("config_snapshot", sa.Text(), nullable=False), + ) + metadata.create_all(engine) + + legacy_soul: dict[str, object] = { + "files": {"files": [], "skills": []}, + "prompt": {"system_prompt": "Build mode"}, + } + with pytest.raises(ValidationError) as exc_info: + AgentSoulConfig.model_validate(legacy_soul) + assert exc_info.value.errors(include_url=False)[0]["loc"] == ("files",) + assert exc_info.value.errors(include_url=False)[0]["type"] == "extra_forbidden" + + with engine.begin() as connection: + for table_name in ("agent_config_snapshots", "agent_config_drafts"): + connection.execute( + sa.text(f"INSERT INTO {table_name} (id, config_snapshot) VALUES (:id, :config_snapshot)"), + {"id": table_name, "config_snapshot": json.dumps(legacy_soul)}, + ) + + _run_upgrade(_load_migration_module(), engine) + + with engine.begin() as connection: + stored_build_draft = connection.execute(sa.text("SELECT config_snapshot FROM agent_config_drafts")).scalar_one() + stored_snapshot = connection.execute(sa.text("SELECT config_snapshot FROM agent_config_snapshots")).scalar_one() + + for stored_soul in (stored_build_draft, stored_snapshot): + value = json.loads(stored_soul) + assert "files" not in value + assert AgentSoulConfig.model_validate(value).prompt.system_prompt == "Build mode" + + +@pytest.mark.parametrize( + ("dialect_name", "removal_expression", "presence_predicate", "other_dialect_expression"), + [ + ( + "postgresql", + "config_snapshot::jsonb - 'files'", + "config_snapshot::jsonb ? 'files'", + "JSON_REMOVE", + ), + ( + "mysql", + "JSON_REMOVE(config_snapshot, '$.files')", + "JSON_CONTAINS_PATH(config_snapshot, 'one', '$.files')", + "config_snapshot::jsonb", + ), + ], +) +def test_upgrade_emits_legacy_files_cleanup_in_offline_sql( + dialect_name: str, + removal_expression: str, + presence_predicate: str, + other_dialect_expression: str, +) -> None: + module = _load_migration_module() + output = StringIO() + migration_context = MigrationContext.configure( + dialect_name=dialect_name, + opts={"as_sql": True, "output_buffer": output}, + ) + operations = Operations(migration_context) + original_op = module.__dict__["op"] + module.__dict__["op"] = operations + try: + module.__dict__["upgrade"]() + finally: + module.__dict__["op"] = original_op + + generated_sql = output.getvalue() + assert "UPDATE agent_config_snapshots" in generated_sql + assert "UPDATE agent_config_drafts" in generated_sql + assert generated_sql.count(removal_expression) == 2 + assert generated_sql.count(presence_predicate) == 2 + assert other_dialect_expression not in generated_sql diff --git a/api/tests/unit_tests/pyrefly.toml b/api/tests/unit_tests/pyrefly.toml index 2b0337fde19..36b9a2e0605 100644 --- a/api/tests/unit_tests/pyrefly.toml +++ b/api/tests/unit_tests/pyrefly.toml @@ -800,11 +800,6 @@ project-excludes = [ "services/rag_pipeline/test_rag_pipeline_service.py", "services/rag_pipeline/test_rag_pipeline_task_proxy.py", "services/rag_pipeline/test_rag_pipeline_transform_service.py", - "services/recommend_app/test_buildin_retrieval.py", - "services/recommend_app/test_category_order.py", - "services/recommend_app/test_recommend_app_factory.py", - "services/recommend_app/test_recommend_app_type.py", - "services/recommend_app/test_remote_retrieval.py", "services/retention/test_messages_clean_policy.py", "services/retention/workflow_run/test_archive_download_preparation.py", "services/retention/workflow_run/test_archive_download_task_cache.py", diff --git a/api/tests/unit_tests/repositories/test_account_repository.py b/api/tests/unit_tests/repositories/test_account_repository.py new file mode 100644 index 00000000000..60c20e9eef3 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_account_repository.py @@ -0,0 +1,72 @@ +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from models.account import Account +from repositories.account_repository import SQLAlchemyAccountRepository +from services.entities.account_entities import AccountProfileChanges + + +def _persist_account(session: Session) -> Account: + account = Account(name="Original", email="account@example.com") + account.id = "account-1" + account.interface_language = "en-US" + account.interface_theme = "light" + account.timezone = "UTC" + session.add(account) + session.commit() + return account + + +def test_update_profile_persists_multiple_fields( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + _persist_account(sqlite_session) + repository = SQLAlchemyAccountRepository(sqlite_session_factory) + + result = repository.update_profile( + "account-1", + AccountProfileChanges( + name="Updated", + avatar="avatar-file", + interface_language="zh-Hans", + interface_theme="dark", + timezone="Asia/Shanghai", + ), + ) + + assert result is not None + assert result.name == "Updated" + sqlite_session.expire_all() + persisted = sqlite_session.get(Account, "account-1") + assert persisted is not None + assert persisted.name == "Updated" + assert persisted.avatar == "avatar-file" + assert persisted.interface_language == "zh-Hans" + assert persisted.interface_theme == "dark" + assert persisted.timezone == "Asia/Shanghai" + + +def test_update_profile_rolls_back_on_error( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + _persist_account(sqlite_session) + repository = SQLAlchemyAccountRepository(sqlite_session_factory) + + def fail_to_create_snapshot(_account: Account) -> None: + raise RuntimeError("abort update") + + monkeypatch.setattr(SQLAlchemyAccountRepository, "_to_snapshot", staticmethod(fail_to_create_snapshot)) + + with pytest.raises(RuntimeError, match="abort update"): + repository.update_profile( + "account-1", + AccountProfileChanges(name="Should Roll Back"), + ) + + sqlite_session.expire_all() + persisted = sqlite_session.get(Account, "account-1") + assert persisted is not None + assert persisted.name == "Original" diff --git a/api/tests/unit_tests/repositories/test_app_definition_query_repository.py b/api/tests/unit_tests/repositories/test_app_definition_query_repository.py index c47943c80df..cc9c851bddf 100644 --- a/api/tests/unit_tests/repositories/test_app_definition_query_repository.py +++ b/api/tests/unit_tests/repositories/test_app_definition_query_repository.py @@ -4,7 +4,7 @@ import pytest from sqlalchemy.orm import Session, sessionmaker from core.tools.entities.tool_entities import ApiProviderSchemaType -from models.account import Account +from models.account import Account, Tenant, TenantStatus from models.enums import CustomizeTokenStrategy, TagType from models.model import App, AppMode, AppModelConfig, IconType, Site, Tag, TagBinding from models.tools import ApiToolProvider @@ -16,6 +16,7 @@ from services.app_definition_query_service import ( AppSiteConfiguration, AppToolIconSource, ) +from services.web_app_runtime_query_service import WebAppRuntimeRecord _APP_ID = "11111111-1111-1111-1111-111111111111" _TENANT_ID = "22222222-2222-2222-2222-222222222222" @@ -341,11 +342,97 @@ def test_get_site_configuration_maps_site_fields(sqlite_session_factory: session input_placeholder="Ask anything", custom_disclaimer="Disclaimer", default_language="en-US", + prompt_public=True, show_workflow_steps=False, use_icon_as_answer_icon=True, ) +@pytest.mark.parametrize( + ("tenant_status", "expected_mode"), + [ + (TenantStatus.NORMAL, AppMode.AGENT_CHAT.value), + (TenantStatus.ARCHIVE, AppMode.CHAT.value), + ], +) +def test_get_runtime_record_maps_app_tenant_site_and_compatible_mode( + sqlite_session_factory: sessionmaker[Session], + tenant_status: TenantStatus, + expected_mode: str, +) -> None: + tenant_custom_config = '{"remove_webapp_brand":true,"replace_webapp_logo":"logo-file"}' + with sqlite_session_factory.begin() as session: + tenant = Tenant( + name="Test Tenant", + plan="pro", + status=tenant_status, + custom_config=tenant_custom_config, + ) + tenant.id = _TENANT_ID + session.add(tenant) + app = _persist_app(session) + app_model_config = AppModelConfig( + app_id=app.id, + agent_mode=json.dumps({"enabled": True, "strategy": "react"}), + ) + session.add(app_model_config) + session.flush() + app.app_model_config_id = app_model_config.id + session.add( + Site( + app_id=app.id, + title="Test Site", + icon_type=IconType.IMAGE, + icon="11111111-1111-4111-8111-111111111111", + icon_background="#ffffff", + default_language="en-US", + chat_color_theme="light", + chat_color_theme_inverted=False, + customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, + prompt_public=True, + show_workflow_steps=True, + use_icon_as_answer_icon=False, + ) + ) + + repository = AppDefinitionQueryRepository(session_factory=sqlite_session_factory) + + assert repository.get_runtime_record(_APP_ID) == WebAppRuntimeRecord( + app_id=_APP_ID, + tenant_id=_TENANT_ID, + mode=expected_mode, + enable_site=True, + site=AppSiteConfiguration( + title="Test Site", + chat_color_theme="light", + chat_color_theme_inverted=False, + icon_type=IconType.IMAGE.value, + icon="11111111-1111-4111-8111-111111111111", + icon_background="#ffffff", + description=None, + copyright=None, + privacy_policy=None, + input_placeholder=None, + custom_disclaimer="", + default_language="en-US", + prompt_public=True, + show_workflow_steps=True, + use_icon_as_answer_icon=False, + ), + plan="pro", + tenant_status=tenant_status.value, + tenant_custom_config_json=tenant_custom_config, + ) + + +def test_get_runtime_record_returns_none_for_missing_app( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = AppDefinitionQueryRepository(session_factory=sqlite_session_factory) + + assert repository.get_runtime_record(_APP_ID) is None + + def _tool(provider_type: str, provider_id: str, tool_name: str) -> dict[str, object]: return { "provider_type": provider_type, diff --git a/api/tests/unit_tests/repositories/test_recommended_app_catalog_repository.py b/api/tests/unit_tests/repositories/test_recommended_app_catalog_repository.py new file mode 100644 index 00000000000..31bf573306f --- /dev/null +++ b/api/tests/unit_tests/repositories/test_recommended_app_catalog_repository.py @@ -0,0 +1,249 @@ +import json +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from sqlalchemy import event +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker + +from extensions.ext_redis import RedisClientWrapper +from models.model import App, AppMode, RecommendedApp, Site +from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository +from services.recommended_app_query_service import RecommendedAppDetailRecord + + +def _add_catalog_app( + session: Session, + *, + categories: list[str] | None = None, + language: str = "en-US", + is_public: bool = True, + is_listed: bool = True, + is_learn_dify: bool = False, + with_site: bool = True, +) -> App: + app = App( + id=str(uuid4()), + tenant_id=str(uuid4()), + name="Recommended App", + mode=AppMode.CHAT, + icon_type=None, + icon=None, + icon_background="#fff", + enable_site=True, + enable_api=True, + is_public=is_public, + ) + recommended_app = RecommendedApp( + app_id=app.id, + description={}, + copyright="copyright", + privacy_policy="privacy", + category="Workflow", + categories=["Workflow"] if categories is None else categories, + custom_disclaimer="catalog disclaimer", + position=1, + is_listed=is_listed, + is_learn_dify=is_learn_dify, + language=language, + ) + session.add_all([app, recommended_app]) + if with_site: + session.add( + Site( + app_id=app.id, + title="Recommended App", + description="site description", + copyright="site copyright", + privacy_policy="site privacy", + custom_disclaimer="site disclaimer", + default_language="en-US", + customize_token_strategy="not_allow", + ) + ) + session.commit() + return app + + +def _redis() -> MagicMock: + redis = MagicMock(spec=RedisClientWrapper) + redis.get.return_value = None + return redis + + +def _repository( + session_factory: sessionmaker[Session], + *, + redis: RedisClientWrapper | None = None, +) -> DatabaseRecommendedAppCatalogRepository: + return DatabaseRecommendedAppCatalogRepository( + session_factory, + redis=redis if redis is not None else _redis(), + ) + + +def test_list_recommended_returns_typed_records_and_falls_back_language( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory() as session: + app = _add_catalog_app(session) + + repository = _repository(sqlite_session_factory) + page = repository.list_recommended("fr-FR") + + assert page.categories == ("Workflow",) + assert len(page.recommended_apps) == 1 + record = page.recommended_apps[0] + assert record.app_id == app.id + assert record.app is not None + assert record.app.id == app.id + assert record.app.mode == "chat" + assert record.description == "site description" + assert record.custom_disclaimer == "site disclaimer" + assert record.categories == ("Workflow",) + + +def test_list_recommended_skips_private_apps_and_apps_without_sites( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory() as session: + _add_catalog_app(session, is_public=False) + _add_catalog_app(session, with_site=False) + + repository = _repository(sqlite_session_factory) + + assert repository.list_recommended("en-US").recommended_apps == () + + +def test_list_recommended_does_not_restore_legacy_category_when_categories_are_empty( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory() as session: + app = _add_catalog_app(session, categories=[]) + + page = _repository(sqlite_session_factory).list_recommended("en-US") + + record = next(item for item in page.recommended_apps if item.app_id == app.id) + assert record.categories == () + assert "Workflow" not in page.categories + + +def test_list_recommended_uses_redis_category_order( + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory() as session: + _add_catalog_app(session, categories=["A", "B", "C", "D"]) + + checked_out_connections = 0 + + def record_checkout(_dbapi_connection, _connection_record, _connection_proxy) -> None: + nonlocal checked_out_connections + checked_out_connections += 1 + + def record_checkin(_dbapi_connection, _connection_record) -> None: + nonlocal checked_out_connections + checked_out_connections -= 1 + + def get_category_order(_key: str) -> bytes: + assert checked_out_connections == 0 + return json.dumps(["C", "A", "B"]).encode() + + redis = _redis() + redis.get.side_effect = get_category_order + event.listen(sqlite_engine, "checkout", record_checkout) + event.listen(sqlite_engine, "checkin", record_checkin) + try: + page = _repository(sqlite_session_factory, redis=redis).list_recommended("en-US") + finally: + event.remove(sqlite_engine, "checkout", record_checkout) + event.remove(sqlite_engine, "checkin", record_checkin) + + assert page.categories == ("C", "A", "B") + redis.get.assert_called_once_with("explore:apps:category_order:en-US") + + +def test_list_recommended_sorts_categories_without_redis_order( + sqlite_session_factory: sessionmaker[Session], +) -> None: + redis = _redis() + with sqlite_session_factory() as session: + _add_catalog_app(session, categories=["B", "A", "C"]) + + page = _repository(sqlite_session_factory, redis=redis).list_recommended("en-US") + + assert page.categories == ("A", "B", "C") + + +def test_list_learn_dify_filters_flag_and_hides_page_categories( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory() as session: + learn_app = _add_catalog_app(session, is_learn_dify=True) + _add_catalog_app(session, is_learn_dify=False) + + redis = _redis() + repository = _repository(sqlite_session_factory, redis=redis) + page = repository.list_learn_dify("fr-FR") + + assert [app.app_id for app in page.recommended_apps] == [learn_app.id] + assert page.recommended_apps[0].categories == ("Workflow",) + assert page.categories == () + redis.get.assert_not_called() + + +def test_membership_does_not_export_dsl( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory() as session: + app = _add_catalog_app(session) + + repository = _repository(sqlite_session_factory) + with patch( + "repositories.recommended_app_catalog_repository.AppDslService.export_dsl", + return_value="exported yaml", + ) as export_dsl: + detail = repository.get_detail(app.id) + is_in_catalog = repository.contains(app.id) + + assert detail == RecommendedAppDetailRecord( + id=app.id, + name="Recommended App", + icon=None, + icon_background="#fff", + mode="chat", + export_data="exported yaml", + ) + assert is_in_catalog is True + export_dsl.assert_called_once() + + +def test_detail_rejects_unlisted_or_private_apps(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory() as session: + private_app = _add_catalog_app(session, is_public=False) + unlisted_app = _add_catalog_app(session, is_listed=False) + missing_app_id = str(uuid4()) + + repository = _repository(sqlite_session_factory) + + assert repository.get_detail(private_app.id) is None + assert repository.get_detail(unlisted_app.id) is None + assert repository.get_detail(missing_app_id) is None + assert repository.contains(private_app.id) is False + assert repository.contains(unlisted_app.id) is False + assert repository.contains(missing_app_id) is False + + +def test_detail_does_not_require_site(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory() as session: + app = _add_catalog_app(session, with_site=False) + + repository = _repository(sqlite_session_factory) + with patch( + "repositories.recommended_app_catalog_repository.AppDslService.export_dsl", + return_value="exported yaml", + ): + detail = repository.get_detail(app.id) + + assert detail is not None + assert detail.id == app.id diff --git a/api/tests/unit_tests/repositories/test_tag_repository.py b/api/tests/unit_tests/repositories/test_tag_repository.py new file mode 100644 index 00000000000..516dc028405 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_tag_repository.py @@ -0,0 +1,129 @@ +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from models.enums import TagType +from models.model import Tag, TagBinding +from models.snippet import CustomizedSnippet, SnippetType +from repositories.tag_repository import TagRepository +from services.tag_application_service import ( + CreateTagInput, + TagBindingInput, + TagBindingTargetNotFoundError, + TagNameConflictError, + TagNotFoundError, + UpdateTagInput, +) + + +def _tag(tag_id: str, *, workspace_id: str, tag_type: TagType, name: str) -> Tag: + tag = Tag(tenant_id=workspace_id, type=tag_type, name=name, created_by="account-1") + tag.id = tag_id + return tag + + +def _snippet(snippet_id: str, *, workspace_id: str) -> CustomizedSnippet: + snippet = CustomizedSnippet( + tenant_id=workspace_id, + name="Snippet", + description="", + type=SnippetType.NODE.value, + created_by="account-1", + updated_by="account-1", + ) + snippet.id = snippet_id + return snippet + + +def test_list_tags_scopes_binding_counts_and_escapes_keyword( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _tag("tag-1", workspace_id="workspace-1", tag_type=TagType.APP, name="50% discount"), + _tag("tag-2", workspace_id="workspace-1", tag_type=TagType.APP, name="500 discount"), + _tag("tag-3", workspace_id="workspace-2", tag_type=TagType.APP, name="50% other"), + TagBinding(tenant_id="workspace-1", tag_id="tag-1", target_id="app-1", created_by="account-1"), + TagBinding(tenant_id="workspace-2", tag_id="tag-1", target_id="app-2", created_by="account-2"), + ] + ) + + result = TagRepository(sqlite_session_factory).list_tags("workspace-1", "app", "50%") + + assert result == (("tag-1", "50% discount", "app", 1),) + + +def test_tag_lifecycle_uses_owned_transactions_and_workspace_scope( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = TagRepository(sqlite_session_factory) + created = repository.create_tag("workspace-1", "account-1", CreateTagInput("Original", "knowledge")) + + with sqlite_session_factory.begin() as session: + session.add( + TagBinding( + tenant_id="workspace-1", + tag_id=created.id, + target_id="dataset-1", + created_by="account-1", + ) + ) + + assert repository.get_tag_type("workspace-1", created.id) == "knowledge" + assert repository.get_tag_type("workspace-2", created.id) is None + + updated = repository.update_tag("workspace-1", created.id, UpdateTagInput("Updated")) + assert updated.name == "Updated" + assert updated.binding_count == 1 + + with pytest.raises(TagNameConflictError): + repository.create_tag("workspace-1", "account-1", CreateTagInput("Updated", "knowledge")) + + with pytest.raises(TagNotFoundError): + repository.update_tag("workspace-2", created.id, UpdateTagInput("Leaked")) + + repository.delete_tag("workspace-1", created.id) + with sqlite_session_factory() as session: + assert session.scalar(select(Tag.id).where(Tag.id == created.id)) is None + assert session.scalar(select(TagBinding.id).where(TagBinding.tag_id == created.id)) is None + + +def test_binding_mutations_validate_target_type_and_workspace( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _snippet("snippet-1", workspace_id="workspace-1"), + _tag("tag-1", workspace_id="workspace-1", tag_type=TagType.SNIPPET, name="Valid"), + _tag("tag-2", workspace_id="workspace-1", tag_type=TagType.APP, name="Wrong type"), + _tag("tag-3", workspace_id="workspace-2", tag_type=TagType.SNIPPET, name="Wrong workspace"), + ] + ) + + repository = TagRepository(sqlite_session_factory) + binding = TagBindingInput(("tag-1", "tag-1", "tag-2", "tag-3"), "snippet-1", "snippet") + repository.create_bindings("workspace-1", "account-1", binding) + repository.create_bindings("workspace-1", "account-1", binding) + + with sqlite_session_factory() as session: + bindings = session.scalars(select(TagBinding).where(TagBinding.target_id == "snippet-1")).all() + assert len(bindings) == 1 + assert bindings[0].tag_id == "tag-1" + assert bindings[0].tenant_id == "workspace-1" + + repository.delete_bindings("workspace-1", binding) + with sqlite_session_factory() as session: + assert session.scalars(select(TagBinding).where(TagBinding.target_id == "snippet-1")).all() == [] + + +def test_binding_mutation_rejects_missing_target(sqlite_session_factory: sessionmaker[Session]) -> None: + repository = TagRepository(sqlite_session_factory) + + with pytest.raises(TagBindingTargetNotFoundError, match="Snippet not found"): + repository.create_bindings( + "workspace-1", + "account-1", + TagBindingInput(("tag-1",), "missing", "snippet"), + ) diff --git a/api/tests/unit_tests/repositories/test_trial_app_query_repository.py b/api/tests/unit_tests/repositories/test_trial_app_query_repository.py new file mode 100644 index 00000000000..0f2bbb3ac46 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_trial_app_query_repository.py @@ -0,0 +1,28 @@ +import uuid +from unittest.mock import MagicMock + +from sqlalchemy.orm import Session, sessionmaker + +from models.model import TrialApp +from repositories.trial_app_query_repository import TrialAppQueryRepository + + +def test_existing_ids_returns_only_trial_apps(sqlite_session_factory: sessionmaker[Session]) -> None: + eligible_id = str(uuid.uuid4()) + other_id = str(uuid.uuid4()) + with sqlite_session_factory() as session: + session.add(TrialApp(app_id=eligible_id, tenant_id=str(uuid.uuid4()))) + session.commit() + + result = TrialAppQueryRepository(sqlite_session_factory).existing_ids([eligible_id, other_id]) + + assert result == frozenset({eligible_id}) + + +def test_existing_ids_skips_session_for_empty_input() -> None: + session_factory = MagicMock(spec=sessionmaker) + + result = TrialAppQueryRepository(session_factory).existing_ids([]) + + assert result == frozenset() + session_factory.assert_not_called() diff --git a/api/tests/unit_tests/repositories/test_trial_app_usage_repository.py b/api/tests/unit_tests/repositories/test_trial_app_usage_repository.py new file mode 100644 index 00000000000..57d7e59f7bd --- /dev/null +++ b/api/tests/unit_tests/repositories/test_trial_app_usage_repository.py @@ -0,0 +1,53 @@ +import uuid + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from models.model import AccountTrialAppRecord +from repositories.trial_app_usage_repository import TrialAppUsageRepository + + +def _record( + session_factory: sessionmaker[Session], + *, + app_id: str, + account_id: str, +) -> AccountTrialAppRecord | None: + with session_factory() as session: + return session.scalar( + select(AccountTrialAppRecord).where( + AccountTrialAppRecord.app_id == app_id, + AccountTrialAppRecord.account_id == account_id, + ) + ) + + +def test_record_increments_existing_usage(sqlite_session_factory: sessionmaker[Session]) -> None: + app_id = str(uuid.uuid4()) + account_id = str(uuid.uuid4()) + with sqlite_session_factory.begin() as session: + session.add(AccountTrialAppRecord(app_id=app_id, account_id=account_id, count=3)) + + TrialAppUsageRepository(sqlite_session_factory).record(app_id=app_id, account_id=account_id) + + record = _record(sqlite_session_factory, app_id=app_id, account_id=account_id) + assert record is not None + assert record.count == 4 + + +def test_record_does_not_commit_caller_session(sqlite_session_factory: sessionmaker[Session]) -> None: + pending_app_id = str(uuid.uuid4()) + usage_app_id = str(uuid.uuid4()) + account_id = str(uuid.uuid4()) + + with sqlite_session_factory() as caller_session: + caller_session.add(AccountTrialAppRecord(app_id=pending_app_id, account_id=account_id, count=1)) + + TrialAppUsageRepository(sqlite_session_factory).record(app_id=usage_app_id, account_id=account_id) + + caller_session.rollback() + + assert _record(sqlite_session_factory, app_id=pending_app_id, account_id=account_id) is None + usage = _record(sqlite_session_factory, app_id=usage_app_id, account_id=account_id) + assert usage is not None + assert usage.count == 1 diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py index 2042dbee794..c3b70798924 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -25,7 +25,6 @@ from models.agent import ( AgentScope, AgentSource, AgentStatus, - AgentWorkspaceBinding, AgentWorkspaceOwnerType, WorkflowAgentBindingType, WorkflowAgentNodeBinding, @@ -562,14 +561,8 @@ def test_save_workflow_composer_commits_before_retiring_replaced_inline_agent( def retire_unowned(**kwargs): assert kwargs["agent_ids"] == {"old-inline-agent"} events.append("retire") - return ["binding-1"], ["home-1"] monkeypatch.setattr(composer_service.WorkflowAgentRetirementService, "retire_unowned", retire_unowned) - monkeypatch.setattr( - composer_service, - "enqueue_agent_resource_collection", - MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), - ) payload = ComposerSavePayload.model_validate( { "variant": ComposerVariant.WORKFLOW, @@ -589,7 +582,7 @@ def test_save_workflow_composer_commits_before_retiring_replaced_inline_agent( payload=payload, ) - assert events == ["commit", "retire", "enqueue"] + assert events == ["commit", "retire"] def test_save_workflow_composer_rejects_agent_app_variant(sqlite_session: Session): @@ -3878,7 +3871,7 @@ def test_reference_counts_include_draft_and_published_bindings_once_per_app(sqli assert result == {"agent-1": 1} -def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): +def test_roster_update_versions_and_detail(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): session = sqlite_session listed_version = AgentConfigSnapshot( id="version-4", @@ -3933,8 +3926,6 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat session.add_all([agent, listed_version, older_listed_version, revision, listed_revision]) session.commit() service = AgentRosterService(session) - retire_snapshots = MagicMock(return_value=[]) - monkeypatch.setattr(AgentHomeSnapshotService, "retire_all_for_agent", retire_snapshots) monkeypatch.setattr( service, "get_roster_agent_detail", @@ -3947,13 +3938,10 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat account_id="account-1", payload=roster_service.RosterAgentUpdatePayload(description="new"), ) - service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") versions = service.list_agent_versions(tenant_id="tenant-1", agent_id="agent-1") detail = service.get_agent_version_detail(tenant_id="tenant-1", agent_id="agent-1", version_id="version-2") assert updated["description"] == "new" - assert agent.status == AgentStatus.ARCHIVED - retire_snapshots.assert_called_once_with(session=session, tenant_id="tenant-1", agent_id="agent-1") assert versions[0]["id"] == "version-4" assert versions[0]["version"] == 2 assert versions[0]["display_version"] == 2 @@ -3970,69 +3958,6 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat assert detail["revisions"][0]["created_at"] == int(revision_created_at.timestamp()) -def test_roster_archive_retires_then_commits_before_enqueue( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - session = sqlite_session - service = AgentRosterService(session) - agent = _agent() - binding = AgentWorkspaceBinding( - id="binding-1", - tenant_id=agent.tenant_id, - app_id="app-1", - workspace_id="workspace-1", - agent_id=agent.id, - agent_config_version_id="version-1", - agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, - backend_binding_ref="backend-binding-1", - ) - session.add_all([agent, binding]) - session.commit() - events: list[str] = [] - monkeypatch.setattr( - AgentWorkspaceService, - "retire_binding", - MagicMock(side_effect=lambda **_kwargs: events.append("retire-binding") or "binding-1"), - ) - monkeypatch.setattr( - AgentHomeSnapshotService, - "retire_all_for_agent", - MagicMock(side_effect=lambda **_kwargs: events.append("retire-home") or ["home-1"]), - ) - event.listen(session, "after_commit", lambda _session: events.append("commit")) - monkeypatch.setattr( - roster_service, - "enqueue_agent_resource_collection", - MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), - ) - - service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") - - assert events == ["retire-binding", "retire-home", "commit", "enqueue"] - - -def test_roster_archive_commit_failure_does_not_enqueue( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - session = sqlite_session - service = AgentRosterService(session) - session.add(_agent()) - session.commit() - monkeypatch.setattr(AgentHomeSnapshotService, "retire_all_for_agent", MagicMock(return_value=["home-1"])) - event.listen( - session, - "before_commit", - lambda _session: (_ for _ in ()).throw(RuntimeError("commit failed")), - ) - enqueue_collection = MagicMock() - monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", enqueue_collection) - - with pytest.raises(RuntimeError, match="commit failed"): - service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") - - enqueue_collection.assert_not_called() - - def test_roster_create_detail_and_lookup_helpers(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): session = sqlite_session service = AgentRosterService(session) diff --git a/api/tests/unit_tests/services/agent/test_deletion_service.py b/api/tests/unit_tests/services/agent/test_deletion_service.py new file mode 100644 index 00000000000..26e9c472a6b --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_deletion_service.py @@ -0,0 +1,373 @@ +from collections.abc import Generator +from contextlib import contextmanager, nullcontext +from decimal import Decimal +from typing import cast +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import Table +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.sql.dml import Delete + +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigRevision, + AgentConfigRevisionOperation, + AgentConfigSnapshot, + AgentConfigVersionKind, + AgentDebugConversation, + AgentHomeSnapshot, + AgentKind, + AgentScope, + AgentSource, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) +from models.agent_config_entities import AgentSoulConfig +from models.enums import AppStatus, ConversationFromSource, ConversationStatus +from models.model import App, AppMode, Conversation, Message +from services.agent.deletion_service import AgentDeletionInvariantError, AgentDeletionService + + +def _archived_agent( + *, + agent_id: str = "agent-1", + tenant_id: str = "tenant-1", + status: AgentStatus = AgentStatus.ARCHIVED, +) -> Agent: + return Agent( + id=agent_id, + tenant_id=tenant_id, + name="Agent", + description="", + role="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=status, + ) + + +def test_purge_archived_agent_deletes_complete_aggregate_and_preserves_workflow_binding( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + agent = _archived_agent() + snapshot = AgentConfigSnapshot( + id="snapshot-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + version=1, + config_snapshot=AgentSoulConfig(), + ) + dangling_binding = WorkflowAgentNodeBinding( + id="workflow-binding-1", + tenant_id=agent.tenant_id, + app_id="missing-app", + workflow_id="missing-workflow", + workflow_version="old-version", + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id=agent.id, + current_snapshot_id=snapshot.id, + node_job_config={}, + ) + rows = [ + agent, + snapshot, + AgentConfigDraft( + id="draft-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + config_snapshot=AgentSoulConfig(), + ), + AgentConfigDraft( + id="build-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + config_snapshot=AgentSoulConfig(), + ), + AgentConfigRevision( + id="revision-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + current_snapshot_id=snapshot.id, + revision=1, + operation=AgentConfigRevisionOperation.CREATE_VERSION, + ), + AgentDebugConversation( + id="debug-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + app_id="app-1", + account_id="account-1", + draft_type=AgentConfigDraftType.DRAFT, + conversation_id="conversation-1", + ), + AgentHomeSnapshot( + id="home-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.RETIRED, + ), + AgentWorkspaceBinding( + id="binding-1", + tenant_id=agent.tenant_id, + app_id="app-1", + workspace_id="workspace-1", + agent_id=agent.id, + agent_config_version_id=snapshot.id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-ref", + status=AgentWorkingResourceStatus.RETIRED, + ), + ] + sibling = _archived_agent(agent_id="agent-2") + other_tenant = _archived_agent(agent_id="agent-3", tenant_id="tenant-2") + unrelated_app = App( + id="unrelated-app", + tenant_id=agent.tenant_id, + name="Unrelated", + mode=AppMode.WORKFLOW, + status=AppStatus.NORMAL, + enable_site=True, + enable_api=True, + ) + conversation = Conversation( + id="conversation-1", + app_id=unrelated_app.id, + mode=AppMode.AGENT_CHAT, + name="Preserved conversation", + _inputs={}, + status=ConversationStatus.NORMAL, + from_source=ConversationFromSource.CONSOLE, + from_account_id="account-1", + ) + preserved_rows = [ + sibling, + other_tenant, + AgentConfigDraft( + id="sibling-draft", + tenant_id=sibling.tenant_id, + agent_id=sibling.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + config_snapshot=AgentSoulConfig(), + ), + AgentHomeSnapshot( + id="other-home", + tenant_id=other_tenant.tenant_id, + agent_id=other_tenant.id, + snapshot_ref="other-home-ref", + status=AgentWorkingResourceStatus.RETIRED, + ), + unrelated_app, + AgentWorkspace( + id="workspace-1", + tenant_id=agent.tenant_id, + app_id=unrelated_app.id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=conversation.id, + owner_scope_key="root", + backend_workspace_ref="workspace-ref", + status=AgentWorkingResourceStatus.ACTIVE, + active_guard=1, + ), + conversation, + Message( + id="message-1", + app_id=unrelated_app.id, + conversation_id=conversation.id, + _inputs={}, + query="hello", + message={"role": "user", "content": "hello"}, + answer="world", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + from_account_id="account-1", + ), + ] + sqlite_session.add_all([*rows, dangling_binding, *preserved_rows]) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + sqlite_session_factory, + ) + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent.id]) + + with sqlite_session_factory() as observer_session: + for row in rows: + assert observer_session.get(type(row), row.id) is None + preserved_binding = observer_session.get(WorkflowAgentNodeBinding, dangling_binding.id) + assert preserved_binding is not None + assert preserved_binding.agent_id == agent.id + for row in preserved_rows: + assert observer_session.get(type(row), row.id) is not None + + +def test_purge_bulk_deletes_aggregate_dependencies_before_agent(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + session.scalars.return_value.all.return_value = [_archived_agent()] + session.scalar.side_effect = [None, None] + deleted_tables: list[str] = [] + + def record_bulk_delete(statement: object) -> None: + if isinstance(statement, Delete): + deleted_tables.append(cast(Table, statement.table).name) + + session.execute.side_effect = record_bulk_delete + monkeypatch.setattr("services.agent.deletion_service.session_factory.create_session", lambda: context) + AgentDeletionService.purge_archived_agents(tenant_id="tenant-1", agent_ids=["agent-1"]) + + assert deleted_tables == [ + cast(Table, model.__table__).name + for model in ( + AgentDebugConversation, + AgentConfigRevision, + AgentConfigDraft, + AgentConfigSnapshot, + AgentHomeSnapshot, + AgentWorkspaceBinding, + Agent, + ) + ] + + +@pytest.mark.parametrize( + ("invariant", "expected_error"), + [ + ("non_archived", "must be ARCHIVED"), + ("active_binding", "still has ACTIVE Binding"), + ("active_home", "still has ACTIVE Home Snapshot"), + ], +) +def test_purge_rejects_invalid_aggregate_invariants( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + invariant: str, + expected_error: str, +) -> None: + agent = _archived_agent(status=AgentStatus.ACTIVE if invariant == "non_archived" else AgentStatus.ARCHIVED) + related: AgentWorkspaceBinding | AgentHomeSnapshot | None = None + if invariant == "active_binding": + related = AgentWorkspaceBinding( + id="binding-1", + tenant_id=agent.tenant_id, + app_id="app-1", + workspace_id="workspace-1", + agent_id=agent.id, + agent_config_version_id="snapshot-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + elif invariant == "active_home": + related = AgentHomeSnapshot( + id="home-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + sqlite_session.add(agent) + if related is not None: + sqlite_session.add(related) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + with pytest.raises(AgentDeletionInvariantError, match=expected_error): + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent.id]) + + assert sqlite_session.get(Agent, agent.id) is not None + if related is not None: + assert sqlite_session.get(type(related), related.id) is not None + + +def test_purge_is_idempotent_for_empty_missing_and_repeated_ids( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + agent = _archived_agent() + sqlite_session.add(agent) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + agent_id = agent.id + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[]) + assert sqlite_session.get(Agent, agent_id) is not None + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=["missing-agent"]) + assert sqlite_session.get(Agent, agent_id) is not None + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent_id, agent_id]) + assert sqlite_session.get(Agent, agent_id) is None + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent_id]) + assert sqlite_session.get(Agent, agent_id) is None + + +@pytest.mark.parametrize("failure_stage", ["delete", "commit"]) +def test_purge_failure_rolls_back_complete_aggregate( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + failure_stage: str, +) -> None: + agent = _archived_agent() + draft = AgentConfigDraft( + id="draft-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + config_snapshot=AgentSoulConfig(), + ) + sqlite_session.add_all([agent, draft]) + sqlite_session.commit() + agent_id = agent.id + draft_id = draft.id + error = RuntimeError(f"{failure_stage} failed") + + @contextmanager + def failing_session() -> Generator[Session]: + with sqlite_session_factory() as service_session: + failure_method = failure_stage if failure_stage == "commit" else "execute" + monkeypatch.setattr(service_session, failure_method, MagicMock(side_effect=error)) + yield service_session + + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + failing_session, + ) + + with pytest.raises(RuntimeError) as exc_info: + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent_id]) + + assert exc_info.value is error + with sqlite_session_factory() as observer_session: + assert observer_session.get(Agent, agent_id) is not None + assert observer_session.get(AgentConfigDraft, draft_id) is not None diff --git a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py index abe565aaa94..b3476f1c120 100644 --- a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py +++ b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py @@ -148,6 +148,42 @@ def test_home_snapshot_collection_database_failure_propagates(monkeypatch: pytes assert exc_info.value is error +@pytest.mark.parametrize("snapshot_state", ["missing", "active"]) +@pytest.mark.parametrize("sqlite_session", [(AgentHomeSnapshot,)], indirect=True) +def test_home_snapshot_collection_non_retired_target_is_idempotent_noop( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + snapshot_state: str, +) -> None: + if snapshot_state == "active": + sqlite_session.add( + AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="snapshot-ref-1", + status=AgentWorkingResourceStatus.ACTIVE, + ) + ) + sqlite_session.commit() + delete = MagicMock() + commit = MagicMock(wraps=sqlite_session.commit) + monkeypatch.setattr( + "services.agent.home_snapshot_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete) + monkeypatch.setattr(sqlite_session, "commit", commit) + + AgentHomeSnapshotService.collect_retired_home_snapshot( + tenant_id="tenant-1", + home_snapshot_id="home-1", + ) + + delete.assert_not_called() + commit.assert_not_called() + + @pytest.mark.parametrize( "sqlite_session", [(AgentHomeSnapshot, AgentConfigDraft, AgentConfigSnapshot)], @@ -185,6 +221,50 @@ def test_home_snapshot_collection_backend_failure_propagates_and_preserves_retir assert stored_snapshot.status is AgentWorkingResourceStatus.RETIRED +@pytest.mark.parametrize( + "sqlite_session", + [(AgentHomeSnapshot, AgentConfigDraft, AgentConfigSnapshot)], + indirect=True, +) +def test_home_snapshot_collection_ignores_config_references( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + snapshot = AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="snapshot-ref-1", + status=AgentWorkingResourceStatus.RETIRED, + ) + draft = _build_draft(home_snapshot_id=snapshot.id) + config_snapshot = AgentConfigSnapshot( + id="config-1", + tenant_id="tenant-1", + agent_id="agent-1", + version=1, + home_snapshot_id=snapshot.id, + config_snapshot=AgentSoulConfig(), + ) + sqlite_session.add_all([snapshot, draft, config_snapshot]) + sqlite_session.commit() + delete = MagicMock() + monkeypatch.setattr( + "services.agent.home_snapshot_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete) + + AgentHomeSnapshotService.collect_retired_home_snapshot( + tenant_id="tenant-1", + home_snapshot_id=snapshot.id, + ) + + delete.assert_called_once_with(snapshot_ref=snapshot.snapshot_ref) + assert sqlite_session.get(AgentHomeSnapshot, snapshot.id) is None + assert sqlite_session.get(AgentConfigDraft, draft.id) is not None + assert sqlite_session.get(AgentConfigSnapshot, config_snapshot.id) is not None + + @pytest.mark.parametrize( "sqlite_session", [(AgentHomeSnapshot, AgentConfigDraft, AgentConfigSnapshot)], diff --git a/api/tests/unit_tests/services/agent/test_retirement_service.py b/api/tests/unit_tests/services/agent/test_retirement_service.py index 71362eecfe3..698f5a2a799 100644 --- a/api/tests/unit_tests/services/agent/test_retirement_service.py +++ b/api/tests/unit_tests/services/agent/test_retirement_service.py @@ -28,10 +28,9 @@ from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.workspace_service import AgentWorkspaceService -def test_retire_unowned_commits_resource_retirement(monkeypatch: pytest.MonkeyPatch) -> None: +def test_retire_unowned_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: context = MagicMock() - session = context.__enter__.return_value - session.scalars.return_value.all.return_value = [SimpleNamespace(id="binding-1")] + error = RuntimeError("retirement failed") monkeypatch.setattr( "services.agent.retirement_service.session_factory.create_session", lambda: context, @@ -39,30 +38,20 @@ def test_retire_unowned_commits_resource_retirement(monkeypatch: pytest.MonkeyPa monkeypatch.setattr( WorkflowAgentRetirementService, "archive_unowned", - MagicMock(return_value=["agent-1"]), - ) - monkeypatch.setattr( - AgentWorkspaceService, - "retire_binding", - MagicMock(return_value="binding-1"), - ) - monkeypatch.setattr( - AgentHomeSnapshotService, - "retire_all_for_agent", - MagicMock(return_value=["home-1"]), + MagicMock(side_effect=error), ) - result = WorkflowAgentRetirementService.retire_unowned( - tenant_id="tenant-1", - agent_ids=["agent-1"], - account_id="account-1", - ) + with pytest.raises(RuntimeError) as exc_info: + WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=["agent-1"], + account_id="account-1", + ) - assert result == (["binding-1"], ["home-1"]) - session.commit.assert_called_once_with() + assert exc_info.value is error -def _workflow_only_agent() -> Agent: +def _workflow_only_agent(*, backing_app_id: str | None = None) -> Agent: return Agent( id="agent-1", tenant_id="tenant-1", @@ -73,17 +62,29 @@ def _workflow_only_agent() -> Agent: scope=AgentScope.WORKFLOW_ONLY, source=AgentSource.WORKFLOW, status=AgentStatus.ACTIVE, + backing_app_id=backing_app_id, ) @pytest.mark.parametrize( - "sqlite_session", - [(Agent, App, Workflow, WorkflowAgentNodeBinding)], - indirect=True, + ("workflow_version", "pointer_to_owner", "mismatched_key", "expected_status"), + [ + pytest.param(Workflow.VERSION_DRAFT, False, None, AgentStatus.ACTIVE, id="draft-owner"), + pytest.param("current-version", True, None, AgentStatus.ACTIVE, id="current-published-owner"), + pytest.param("historical-version", False, None, AgentStatus.ACTIVE, id="historical-published-owner"), + pytest.param("v1", True, "tenant_id", AgentStatus.ARCHIVED, id="tenant-mismatch"), + pytest.param("v1", True, "app_id", AgentStatus.ARCHIVED, id="app-mismatch"), + pytest.param("v1", True, "workflow_id", AgentStatus.ARCHIVED, id="workflow-mismatch"), + pytest.param("v1", True, "workflow_version", AgentStatus.ARCHIVED, id="version-mismatch"), + ], ) -def test_retire_unowned_keeps_effectively_owned_agent_active( +def test_retire_unowned_requires_an_exact_persisted_workflow_owner_key( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, + workflow_version: str, + pointer_to_owner: bool, + mismatched_key: str | None, + expected_status: AgentStatus, ) -> None: agent = _workflow_only_agent() app = App( @@ -96,10 +97,10 @@ def test_retire_unowned_keeps_effectively_owned_agent_active( enable_api=True, ) workflow = Workflow.new( - tenant_id="tenant-1", + tenant_id="workflow-tenant" if mismatched_key == "tenant_id" else "tenant-1", app_id=app.id, type=WorkflowType.WORKFLOW.value, - version=Workflow.VERSION_DRAFT, + version=workflow_version, graph="{}", features="{}", created_by="account-1", @@ -107,11 +108,22 @@ def test_retire_unowned_keeps_effectively_owned_agent_active( conversation_variables=[], rag_pipeline_variables=[], ) + app.workflow_id = workflow.id if pointer_to_owner else "another-current-workflow" + binding_key = { + "tenant_id": "tenant-1", + "app_id": workflow.app_id, + "workflow_id": workflow.id, + "workflow_version": workflow.version, + } + mismatched_values = { + "app_id": "app-2", + "workflow_id": "workflow-2", + "workflow_version": "other-version", + } + if mismatched_key is not None and mismatched_key != "tenant_id": + binding_key[mismatched_key] = mismatched_values[mismatched_key] binding = WorkflowAgentNodeBinding( - tenant_id="tenant-1", - app_id=app.id, - workflow_id=workflow.id, - workflow_version=workflow.version, + **binding_key, node_id="agent-node", binding_type=WorkflowAgentBindingType.INLINE_AGENT, agent_id=agent.id, @@ -124,17 +136,21 @@ def test_retire_unowned_keeps_effectively_owned_agent_active( "services.agent.retirement_service.session_factory.create_session", lambda: nullcontext(sqlite_session), ) - - result = WorkflowAgentRetirementService.retire_unowned( + celery_delay = MagicMock() + monkeypatch.setattr("tasks.collect_agent_resources_task.collect_agent_resources.delay", celery_delay) + WorkflowAgentRetirementService.retire_unowned( tenant_id="tenant-1", agent_ids=[agent.id], account_id="account-1", ) - assert result == ([], []) stored_agent = sqlite_session.get(Agent, agent.id) assert stored_agent is not None - assert stored_agent.status is AgentStatus.ACTIVE + assert stored_agent.status is expected_status + if expected_status is AgentStatus.ACTIVE: + celery_delay.assert_not_called() + else: + celery_delay.assert_called_once() @pytest.mark.parametrize( @@ -146,7 +162,16 @@ def test_retire_unowned_archives_orphan_and_retires_resources( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, ) -> None: - agent = _workflow_only_agent() + agent = _workflow_only_agent(backing_app_id="hidden-app-1") + hidden_app = App( + id="hidden-app-1", + tenant_id="tenant-1", + name="Inline Agent runtime", + mode=AppMode.AGENT, + status=AppStatus.NORMAL, + enable_site=True, + enable_api=True, + ) home = AgentHomeSnapshot( id="home-1", tenant_id="tenant-1", @@ -157,7 +182,7 @@ def test_retire_unowned_archives_orphan_and_retires_resources( workspace = AgentWorkspace( id="workspace-1", tenant_id="tenant-1", - app_id="app-1", + app_id=hidden_app.id, owner_type=AgentWorkspaceOwnerType.CONVERSATION, owner_id="conversation-1", owner_scope_key="root", @@ -168,7 +193,7 @@ def test_retire_unowned_archives_orphan_and_retires_resources( binding = AgentWorkspaceBinding( id="binding-1", tenant_id="tenant-1", - app_id="app-1", + app_id=hidden_app.id, workspace_id=workspace.id, agent_id=agent.id, base_home_snapshot_id=home.id, @@ -177,20 +202,26 @@ def test_retire_unowned_archives_orphan_and_retires_resources( backend_binding_ref="binding-ref", status=AgentWorkingResourceStatus.ACTIVE, ) - sqlite_session.add_all([agent, home, workspace, binding]) + sqlite_session.add_all([agent, hidden_app, home, workspace, binding]) sqlite_session.commit() monkeypatch.setattr( "services.agent.retirement_service.session_factory.create_session", lambda: nullcontext(sqlite_session), ) + cleanup_app = MagicMock() + enqueue_collection = MagicMock() + monkeypatch.setattr("services.agent.retirement_service.remove_app_and_related_data_task.delay", cleanup_app) + monkeypatch.setattr( + "services.agent.retirement_service.enqueue_agent_resource_collection", + enqueue_collection, + ) - result = WorkflowAgentRetirementService.retire_unowned( + WorkflowAgentRetirementService.retire_unowned( tenant_id="tenant-1", agent_ids=[agent.id], account_id="account-1", ) - assert result == ([binding.id], [home.id]) stored_agent = sqlite_session.get(Agent, agent.id) stored_binding = sqlite_session.get(AgentWorkspaceBinding, binding.id) stored_workspace = sqlite_session.get(AgentWorkspace, workspace.id) @@ -199,7 +230,165 @@ def test_retire_unowned_archives_orphan_and_retires_resources( assert stored_binding is not None assert stored_workspace is not None assert stored_home is not None + assert sqlite_session.get(App, hidden_app.id) is None assert stored_agent.status is AgentStatus.ARCHIVED assert stored_binding.status is AgentWorkingResourceStatus.RETIRED assert stored_workspace.status is AgentWorkingResourceStatus.RETIRED assert stored_home.status is AgentWorkingResourceStatus.RETIRED + cleanup_app.assert_called_once_with(tenant_id="tenant-1", app_id=hidden_app.id) + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + workspace_ids=[workspace.id], + binding_ids=[binding.id], + home_snapshot_ids=[home.id], + purge_agent_ids=[agent.id], + ) + + +def test_hidden_app_enqueue_failure_prevents_agent_purge_enqueue(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + session.scalars.side_effect = [ + SimpleNamespace( + all=MagicMock( + return_value=[ + SimpleNamespace(backing_app_id="hidden-app-1"), + SimpleNamespace(backing_app_id="hidden-app-2"), + ] + ) + ), + SimpleNamespace(all=MagicMock(return_value=[])), + SimpleNamespace(all=MagicMock(return_value=[])), + SimpleNamespace(all=MagicMock(return_value=[])), + SimpleNamespace(all=MagicMock(return_value=[])), + ] + monkeypatch.setattr( + "services.agent.retirement_service.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr( + WorkflowAgentRetirementService, + "archive_unowned", + MagicMock(return_value=["agent-1", "agent-2"]), + ) + monkeypatch.setattr(AgentWorkspaceService, "retire_all_for_app", MagicMock(return_value=[])) + monkeypatch.setattr(AgentHomeSnapshotService, "retire_all_for_agent", MagicMock(return_value=[])) + error = RuntimeError("broker unavailable") + cleanup_app = MagicMock(side_effect=[None, error]) + monkeypatch.setattr("services.agent.retirement_service.remove_app_and_related_data_task.delay", cleanup_app) + enqueue_collection = MagicMock() + monkeypatch.setattr( + "services.agent.retirement_service.enqueue_agent_resource_collection", + enqueue_collection, + ) + + with pytest.raises(RuntimeError) as exc_info: + WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=["agent-1", "agent-2"], + account_id="account-1", + ) + + assert exc_info.value is error + assert [call.kwargs["app_id"] for call in cleanup_app.call_args_list] == ["hidden-app-1", "hidden-app-2"] + enqueue_collection.assert_not_called() + + +def test_retire_unowned_retry_after_hidden_app_enqueue_failure_preserves_full_collector_payload( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + agent = _workflow_only_agent(backing_app_id="hidden-app-1") + hidden_app = App( + id="hidden-app-1", + tenant_id="tenant-1", + name="Inline Agent runtime", + mode=AppMode.AGENT, + status=AppStatus.NORMAL, + enable_site=False, + enable_api=False, + ) + home = AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id=agent.id, + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + workspace = AgentWorkspace( + id="workspace-1", + tenant_id="tenant-1", + app_id=hidden_app.id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id="conversation-1", + owner_scope_key="root", + backend_workspace_ref="workspace-ref", + status=AgentWorkingResourceStatus.ACTIVE, + active_guard=1, + ) + binding = AgentWorkspaceBinding( + id="binding-1", + tenant_id="tenant-1", + app_id=hidden_app.id, + workspace_id=workspace.id, + agent_id=agent.id, + base_home_snapshot_id=home.id, + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + sqlite_session.add_all([agent, hidden_app, home, workspace, binding]) + sqlite_session.commit() + agent_id = agent.id + hidden_app_id = hidden_app.id + home_id = home.id + workspace_id = workspace.id + binding_id = binding.id + error = RuntimeError("broker unavailable") + cleanup_app = MagicMock(side_effect=[error, None]) + enqueue_collection = MagicMock() + monkeypatch.setattr("services.agent.retirement_service.remove_app_and_related_data_task.delay", cleanup_app) + monkeypatch.setattr( + "services.agent.retirement_service.enqueue_agent_resource_collection", + enqueue_collection, + ) + + with pytest.raises(RuntimeError) as exc_info: + WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=[agent_id], + account_id="account-1", + ) + + assert exc_info.value is error + sqlite_session.expire_all() + stored_agent = sqlite_session.get(Agent, agent_id) + stored_workspace = sqlite_session.get(AgentWorkspace, workspace_id) + stored_binding = sqlite_session.get(AgentWorkspaceBinding, binding_id) + stored_home = sqlite_session.get(AgentHomeSnapshot, home_id) + assert stored_agent is not None + assert stored_workspace is not None + assert stored_binding is not None + assert stored_home is not None + assert stored_agent.status is AgentStatus.ARCHIVED + assert sqlite_session.get(App, hidden_app_id) is None + assert stored_workspace.status is AgentWorkingResourceStatus.RETIRED + assert stored_binding.status is AgentWorkingResourceStatus.RETIRED + assert stored_home.status is AgentWorkingResourceStatus.RETIRED + enqueue_collection.assert_not_called() + + WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=[agent_id], + account_id="account-1", + ) + + assert cleanup_app.call_count == 2 + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + workspace_ids=[workspace_id], + binding_ids=[binding_id], + home_snapshot_ids=[home_id], + purge_agent_ids=[agent_id], + ) diff --git a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py index 87f376e1c4b..66f06b5ab78 100644 --- a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py +++ b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py @@ -5,7 +5,7 @@ import pytest from sqlalchemy import select from sqlalchemy.orm import Session -from models.agent import Agent, WorkflowAgentBindingType, WorkflowAgentNodeBinding +from models.agent import Agent, AgentScope, WorkflowAgentBindingType, WorkflowAgentNodeBinding from models.enums import AppStatus from models.model import App, AppMode from models.workflow import Workflow, WorkflowType @@ -66,6 +66,38 @@ def test_inline_binding_from_another_node_is_cloned(monkeypatch: pytest.MonkeyPa assert binding.node_job_config.workflow_prompt == "Summarize the input" +def test_draft_sync_resolves_roster_agents() -> None: + draft_workflow = _workflow() + draft_workflow.graph = ( + '{"nodes":[' + '{"id":"node-b","data":{"type":"agent","version":"2","agent_node_kind":"dify_agent",' + '"agent_binding":{"binding_type":"roster_agent","agent_id":"agent-b"}}},' + '{"id":"node-a","data":{"type":"agent","version":"2","agent_node_kind":"dify_agent",' + '"agent_binding":{"binding_type":"roster_agent","agent_id":"agent-a"}}}' + '],"edges":[]}' + ) + session = Mock() + session.scalars.return_value = SimpleNamespace(all=lambda: []) + agents = { + agent_id: SimpleNamespace( + id=agent_id, + scope=AgentScope.ROSTER, + active_config_snapshot_id=f"{agent_id}-snapshot", + ) + for agent_id in ("agent-a", "agent-b") + } + session.scalar.side_effect = [agents["agent-b"], agents["agent-a"]] + + WorkflowAgentPublishService.sync_agent_bindings_for_draft( + session=session, + draft_workflow=draft_workflow, + account_id="account-1", + ) + + assert session.scalar.call_count == 2 + assert {call.args[0].agent_id for call in session.add.call_args_list} == {"agent-a", "agent-b"} + + def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> None: existing_inline = WorkflowAgentNodeBinding( tenant_id="tenant-1", @@ -108,6 +140,11 @@ def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> N SimpleNamespace(all=lambda: [existing_inline, existing_roster]), SimpleNamespace(all=lambda: [source]), ] + session.scalar.return_value = SimpleNamespace( + id="roster-agent", + scope=AgentScope.ROSTER, + active_config_snapshot_id="published-snapshot", + ) retirement_candidates = WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( session=session, source_workflow=_workflow(workflow_id="published-workflow", version="2026-07-13 00:00:00"), @@ -126,16 +163,54 @@ def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> N assert restored.agent_id == "roster-agent" assert restored.current_snapshot_id == "published-snapshot" assert restored.node_job_config.workflow_prompt == "Use the roster agent" - session.flush.assert_called_once() assert retirement_candidates == {"old-inline-agent"} +def test_publish_copy_uses_current_roster_snapshot() -> None: + draft_workflow = _workflow() + draft_workflow.graph = ( + '{"nodes":[{"id":"agent-node","data":{"type":"agent","version":"2",' + '"agent_node_kind":"dify_agent"}}],"edges":[]}' + ) + published_workflow = _workflow(workflow_id="published", version="published") + binding = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version=Workflow.VERSION_DRAFT, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="roster-agent", + current_snapshot_id="old-snapshot", + node_job_config={}, + created_by="account-1", + ) + session = Mock() + active_agent = SimpleNamespace( + id="roster-agent", + scope=AgentScope.ROSTER, + active_config_snapshot_id="active-snapshot", + ) + session.scalar.return_value = active_agent + session.scalars.return_value = SimpleNamespace(all=lambda: [binding]) + + WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + session=session, + draft_workflow=draft_workflow, + published_workflow=published_workflow, + ) + + copied = session.add.call_args.args[0] + assert copied.agent_id == "roster-agent" + assert copied.current_snapshot_id == "active-snapshot" + + @pytest.mark.parametrize( "sqlite_session", [(App, Agent, WorkflowAgentNodeBinding)], indirect=True, ) -def test_publish_binding_replacement_returns_only_previous_inline_agent( +def test_publish_binding_copy_keeps_previous_published_owner( sqlite_session: Session, ) -> None: draft_workflow = _workflow() @@ -198,13 +273,15 @@ def test_publish_binding_replacement_returns_only_previous_inline_agent( ) sqlite_session.add_all([app, previous_inline_binding, previous_roster_binding, draft_binding]) sqlite_session.commit() - retirement_candidates = WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + result = WorkflowAgentPublishService.copy_agent_node_bindings_to_published( session=sqlite_session, draft_workflow=draft_workflow, published_workflow=published_workflow, ) - assert retirement_candidates == {"previous-inline-agent"} + assert result is None + assert sqlite_session.get(WorkflowAgentNodeBinding, previous_inline_binding.id) is previous_inline_binding + assert sqlite_session.get(WorkflowAgentNodeBinding, previous_roster_binding.id) is previous_roster_binding copied = sqlite_session.scalar( select(WorkflowAgentNodeBinding).where( WorkflowAgentNodeBinding.workflow_id == "published-new", @@ -306,7 +383,6 @@ def test_resolve_existing_inline_binding_agent_returns_valid_agent_or_none(monke def test_resolve_roster_binding_rejects_unpublished_agent() -> None: session = Mock() session.scalar.return_value = None - with pytest.raises(ValueError, match="unavailable or unpublished roster agent"): WorkflowAgentPublishService._resolve_roster_agent_graph_binding( session=session, diff --git a/api/tests/unit_tests/services/agent/test_workspace_service.py b/api/tests/unit_tests/services/agent/test_workspace_service.py index d678a92106d..8191f04cbfc 100644 --- a/api/tests/unit_tests/services/agent/test_workspace_service.py +++ b/api/tests/unit_tests/services/agent/test_workspace_service.py @@ -15,7 +15,12 @@ from models.agent import ( AgentWorkspaceBinding, AgentWorkspaceOwnerType, ) -from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope +from services.agent.workspace_service import ( + AgentWorkspaceError, + AgentWorkspaceNotFoundError, + AgentWorkspaceService, + WorkspaceOwnerScope, +) def _scope() -> WorkspaceOwnerScope: @@ -412,9 +417,11 @@ def test_collect_workspace_destroys_workspace_then_remaining_bindings( sqlite_session.add_all([workspace, anchor, remaining]) sqlite_session.commit() client = MagicMock() + commit = MagicMock(wraps=sqlite_session.commit) monkeypatch.setattr( "services.agent.workspace_service.session_factory.create_session", lambda: nullcontext(sqlite_session) ) + monkeypatch.setattr(sqlite_session, "commit", commit) monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id=workspace.id) @@ -429,6 +436,73 @@ def test_collect_workspace_destroys_workspace_then_remaining_bindings( assert sqlite_session.get(AgentWorkspace, workspace.id) is None assert sqlite_session.get(AgentWorkspaceBinding, anchor.id) is None assert sqlite_session.get(AgentWorkspaceBinding, remaining.id) is None + commit.assert_called_once() + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_collect_workspace_remaining_failure_preserves_ledgers_and_replay_converges( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace(status=AgentWorkingResourceStatus.RETIRED) + anchor = _binding(status=AgentWorkingResourceStatus.RETIRED) + remaining = [ + _binding( + binding_id=f"binding-{index}", + agent_id=f"agent-{index}", + status=AgentWorkingResourceStatus.RETIRED, + ) + for index in (2, 3) + ] + anchor.created_at = datetime(2026, 7, 23, 10) + for index, binding in enumerate(remaining, start=1): + binding.created_at = anchor.created_at + timedelta(minutes=index) + sqlite_session.add_all([workspace, anchor, *remaining]) + sqlite_session.commit() + workspace_id = workspace.id + binding_ids = [anchor.id, *(binding.id for binding in remaining)] + error = RuntimeError("middle Binding destroy failed") + client = MagicMock() + client.destroy_execution_binding_sync.side_effect = [None, error, None] + monkeypatch.setattr( + "services.agent.workspace_service.session_factory.create_session", lambda: nullcontext(sqlite_session) + ) + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + + with pytest.raises(RuntimeError) as exc_info: + AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id=workspace_id) + + assert exc_info.value is error + assert client.destroy_execution_binding_sync.call_count == 3 + first_attempt = [call.args[0] for call in client.destroy_execution_binding_sync.call_args_list] + assert [request.destroy_workspace for request in first_attempt] == [True, False, False] + assert sqlite_session.get(AgentWorkspace, workspace_id) is not None + assert all(sqlite_session.get(AgentWorkspaceBinding, binding_id) is not None for binding_id in binding_ids) + + client.reset_mock() + client.destroy_execution_binding_sync.side_effect = None + AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id=workspace_id) + + assert client.destroy_execution_binding_sync.call_count == 3 + assert sqlite_session.get(AgentWorkspace, workspace_id) is None + assert all(sqlite_session.get(AgentWorkspaceBinding, binding_id) is None for binding_id in binding_ids) + + +def test_collect_retired_workspace_without_retired_binding_raises( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace(status=AgentWorkingResourceStatus.RETIRED) + sqlite_session.add(workspace) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.workspace_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + with pytest.raises(AgentWorkspaceError, match="tenant_id=tenant-1, workspace_id=workspace-1"): + AgentWorkspaceService.collect_retired_workspace( + tenant_id="tenant-1", + workspace_id=workspace.id, + ) def test_binding_collection_database_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/services/data_migration/test_import_service.py b/api/tests/unit_tests/services/data_migration/test_import_service.py index 02a897a8f20..80e155c1dae 100644 --- a/api/tests/unit_tests/services/data_migration/test_import_service.py +++ b/api/tests/unit_tests/services/data_migration/test_import_service.py @@ -1,5 +1,8 @@ +from contextlib import nullcontext from dataclasses import dataclass +from types import SimpleNamespace from typing import cast +from unittest.mock import Mock import pytest import yaml @@ -332,7 +335,7 @@ def test_find_existing_mcp_tool_does_not_compare_invalid_uuid(database: Database assert f"{MCPToolProvider.__tablename__}.name" not in where_clause -def test_workflow_app_import_does_not_wrap_app_dsl_import_in_nested_transaction( +def test_workflow_app_import_closes_read_transaction_before_dsl_overwrite( monkeypatch: pytest.MonkeyPatch, database: Database ): class StubAppDslService: @@ -340,32 +343,25 @@ def test_workflow_app_import_does_not_wrap_app_dsl_import_in_nested_transaction( self.session = session def import_app(self, **kwargs): + assert not self.session.in_transaction() return Import(id="import-id", status=ImportStatus.COMPLETED, app_id="imported-app-id") monkeypatch.setattr(import_service, "AppDslService", StubAppDslService) - nested_transactions = [] + monkeypatch.setattr(import_service.dify_config, "RBAC_ENABLED", True) + existing_app = _persist_app(database.session, app_id="11111111-1111-4111-8111-111111111111") + database.session.begin() - def capture_transaction(_session, transaction) -> None: - if transaction.nested: - nested_transactions.append(transaction) - - event.listen(database.session, "after_transaction_create", capture_transaction) - - try: - imported_app_id = MigrationImportService()._import_workflow_app( - account=object(), - workflow_data={"name": "main_chatflow"}, - dsl_content="app:\n mode: workflow\n", - app_id="source-app-id", - existing_app=None, - options=ImportOptions(id_strategy=IdStrategy.PRESERVE_ID), - session=database.session, - ) - finally: - event.remove(database.session, "after_transaction_create", capture_transaction) + imported_app_id = MigrationImportService()._import_workflow_app( + account=object(), + workflow_data={"name": "main_chatflow"}, + dsl_content="app:\n mode: workflow\n", + app_id="source-app-id", + existing_app=existing_app, + options=ImportOptions(id_strategy=IdStrategy.PRESERVE_ID), + session=database.session, + ) assert imported_app_id == "imported-app-id" - assert nested_transactions == [] def test_rewrite_workflow_dsl_replaces_tool_provider_ids(): @@ -489,6 +485,30 @@ def test_workflow_tool_import_publishes_referenced_app_before_create( assert events == [("published", app_id), ("created", app_id)] +def test_ensure_workflow_app_is_published_updates_current_workflow( + monkeypatch: pytest.MonkeyPatch, + database: Database, +) -> None: + _, account = _persist_tenant_account(database.session) + app_id = "00000000-0000-0000-0000-000000000001" + _persist_app(database.session, app_id=app_id) + publish = Mock(return_value=SimpleNamespace(id="published-workflow")) + monkeypatch.setattr(import_service, "WorkflowService", Mock(return_value=SimpleNamespace(publish_workflow=publish))) + monkeypatch.setattr( + import_service, + "sessionmaker", + lambda _engine: SimpleNamespace(begin=lambda: nullcontext(database.session)), + ) + MigrationImportService()._ensure_workflow_app_is_published( + ImportTarget("tenant-1", "target", "account-1", "owner@example.com"), + account, + app_id, + session=database.session, + ) + + assert database.session.get(App, app_id).workflow_id == "published-workflow" + + @pytest.mark.parametrize("id_strategy", [IdStrategy.PRESERVE_ID, IdStrategy.GENERATE_NEW_ID]) def test_workflow_tool_import_id_follows_id_strategy( monkeypatch: pytest.MonkeyPatch, database: Database, id_strategy: IdStrategy diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py index 831bb49dc90..cea5bc86a02 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import pytest from pytest_mock import MockerFixture +from sqlalchemy.dialects import postgresql from sqlalchemy.orm import Session, sessionmaker from core.app.entities.app_invoke_entities import InvokeFrom @@ -69,6 +70,22 @@ class MockRepo: pass +def test_get_published_workflow_by_id_locks_restore_source(mocker: MockerFixture) -> None: + session = mocker.Mock(spec=Session) + workflow = _make_workflow() + workflow.version = "v1" + session.scalar.return_value = workflow + service = RagPipelineService.__new__(RagPipelineService) + service._session = session + + result = service.get_published_workflow_by_id(_make_pipeline(), workflow.id) + + stmt = session.scalar.call_args.args[0] + sql = str(stmt.compile(dialect=postgresql.dialect())) + assert result is workflow + assert "FOR UPDATE" in sql + + def _make_account(account_id: str = "u1", tenant_id: str = "t1") -> Account: account = Account(name="Test User", email=f"{account_id}@example.com") account.id = account_id diff --git a/api/tests/unit_tests/services/recommend_app/__init__.py b/api/tests/unit_tests/services/recommend_app/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/api/tests/unit_tests/services/recommend_app/test_buildin_retrieval.py b/api/tests/unit_tests/services/recommend_app/test_buildin_retrieval.py deleted file mode 100644 index 638deea417e..00000000000 --- a/api/tests/unit_tests/services/recommend_app/test_buildin_retrieval.py +++ /dev/null @@ -1,138 +0,0 @@ -import json -from pathlib import Path -from unittest.mock import MagicMock, patch - -import pytest -import yaml -from sqlalchemy.orm import Session - -from services.recommend_app.buildin.buildin_retrieval import BuildInRecommendAppRetrieval -from services.recommend_app.recommend_app_type import RecommendAppType - -SAMPLE_BUILTIN_DATA = { - "recommended_apps": { - "en-US": {"categories": ["writing"], "apps": [{"id": "app-1"}]}, - "zh-Hans": {"categories": ["search"], "apps": [{"id": "app-2"}]}, - }, - "app_details": { - "app-1": {"id": "app-1", "name": "Writer", "mode": "chat"}, - "app-2": {"id": "app-2", "name": "Searcher", "mode": "workflow"}, - }, -} - - -@pytest.fixture(autouse=True) -def _reset_cache(): - BuildInRecommendAppRetrieval.builtin_data = None - yield - BuildInRecommendAppRetrieval.builtin_data = None - - -class TestBuildInRecommendAppRetrieval: - def test_get_type(self): - retrieval = BuildInRecommendAppRetrieval() - assert retrieval.get_type() == RecommendAppType.BUILDIN - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_get_recommended_apps_and_categories_delegates(self, sqlite_session: Session): - with patch.object( - BuildInRecommendAppRetrieval, - "fetch_recommended_apps_from_builtin", - return_value={"apps": []}, - ) as mock_fetch: - retrieval = BuildInRecommendAppRetrieval() - result = retrieval.get_recommended_apps_and_categories("en-US", session=sqlite_session) - mock_fetch.assert_called_once_with("en-US") - assert result == {"apps": []} - assert not sqlite_session.in_transaction() - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_get_learn_dify_apps_delegates_to_database(self, sqlite_session: Session): - expected = {"recommended_apps": [{"id": "learn-dify-app"}]} - with patch( - "services.recommend_app.buildin.buildin_retrieval.DatabaseRecommendAppRetrieval" - ) as mock_database_retrieval: - mock_database_retrieval.fetch_learn_dify_apps_from_db.return_value = expected - - result = BuildInRecommendAppRetrieval().get_learn_dify_apps("en-US", session=sqlite_session) - - assert result == expected - mock_database_retrieval.fetch_learn_dify_apps_from_db.assert_called_once_with("en-US", session=sqlite_session) - assert not sqlite_session.in_transaction() - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_get_recommend_app_detail_delegates(self, sqlite_session: Session): - with patch.object( - BuildInRecommendAppRetrieval, - "fetch_recommended_app_detail_from_builtin", - return_value={"id": "app-1"}, - ) as mock_fetch: - retrieval = BuildInRecommendAppRetrieval() - result = retrieval.get_recommend_app_detail("app-1", session=sqlite_session) - mock_fetch.assert_called_once_with("app-1") - assert result == {"id": "app-1"} - assert not sqlite_session.in_transaction() - - def test_get_builtin_data_reads_json_and_caches(self, tmp_path: Path): - json_file = tmp_path / "constants" / "recommended_apps.json" - json_file.parent.mkdir(parents=True) - json_file.write_text(json.dumps(SAMPLE_BUILTIN_DATA)) - - mock_app = MagicMock() - mock_app.root_path = str(tmp_path) - - with patch( - "services.recommend_app.buildin.buildin_retrieval.current_app", - mock_app, - ): - first = BuildInRecommendAppRetrieval._get_builtin_data() - second = BuildInRecommendAppRetrieval._get_builtin_data() - - assert first == SAMPLE_BUILTIN_DATA - assert first is second - - def test_fetch_recommended_apps_from_builtin(self): - BuildInRecommendAppRetrieval.builtin_data = SAMPLE_BUILTIN_DATA - result = BuildInRecommendAppRetrieval.fetch_recommended_apps_from_builtin("en-US") - assert result == SAMPLE_BUILTIN_DATA["recommended_apps"]["en-US"] - - def test_fetch_recommended_apps_from_builtin_missing_language(self): - BuildInRecommendAppRetrieval.builtin_data = SAMPLE_BUILTIN_DATA - result = BuildInRecommendAppRetrieval.fetch_recommended_apps_from_builtin("fr-FR") - assert result == {} - - def test_fetch_recommended_app_detail_from_builtin(self): - BuildInRecommendAppRetrieval.builtin_data = SAMPLE_BUILTIN_DATA - result = BuildInRecommendAppRetrieval.fetch_recommended_app_detail_from_builtin("app-1") - assert result == {"id": "app-1", "name": "Writer", "mode": "chat"} - - def test_fetch_recommended_app_detail_from_builtin_missing(self): - BuildInRecommendAppRetrieval.builtin_data = SAMPLE_BUILTIN_DATA - result = BuildInRecommendAppRetrieval.fetch_recommended_app_detail_from_builtin("nonexistent") - assert result is None - - -def test_builtin_workflow_templates_have_unique_end_output_variables(): - """Workflow publish validation rejects duplicate End output variable names, so the bundled - templates must not ship with duplicates or users cannot publish them (see issue #38278).""" - data_path = Path(__file__).resolve().parents[4] / "constants" / "recommended_apps.json" - data = json.loads(data_path.read_text(encoding="utf-8")) - - offenders: dict[str, list[str]] = {} - for app_id, detail in data.get("app_details", {}).items(): - export_data = detail.get("export_data") - if not export_data: - continue - dsl = yaml.safe_load(export_data) - nodes = (dsl or {}).get("workflow", {}).get("graph", {}).get("nodes", []) - output_names = [ - output.get("variable") - for node in nodes - if node.get("data", {}).get("type") == "end" - for output in (node.get("data", {}).get("outputs") or []) - ] - duplicates = sorted({name for name in output_names if output_names.count(name) > 1}) - if duplicates: - offenders[detail.get("name", app_id).strip()] = duplicates - - assert offenders == {}, f"templates with duplicate End output variable names: {offenders}" diff --git a/api/tests/unit_tests/services/recommend_app/test_category_order.py b/api/tests/unit_tests/services/recommend_app/test_category_order.py deleted file mode 100644 index 3b94021f26a..00000000000 --- a/api/tests/unit_tests/services/recommend_app/test_category_order.py +++ /dev/null @@ -1,26 +0,0 @@ -import json -from unittest.mock import patch - -from services.recommend_app.category_order import get_explore_app_category_order, order_categories - - -@patch("services.recommend_app.category_order.redis_client.get") -def test_get_explore_app_category_order_returns_redis_list(mock_get): - mock_get.return_value = json.dumps(["C", "A", "B"]).encode() - - assert get_explore_app_category_order("en-US") == ["C", "A", "B"] - mock_get.assert_called_once_with("explore:apps:category_order:en-US") - - -@patch("services.recommend_app.category_order.redis_client.get") -def test_order_categories_uses_redis_order_as_source_of_truth(mock_get): - mock_get.return_value = json.dumps(["C", "A", "B"]).encode() - - assert order_categories({"A", "B", "C", "D"}, "en-US") == ["C", "A", "B"] - - -@patch("services.recommend_app.category_order.redis_client.get") -def test_order_categories_falls_back_to_sorted_categories_without_redis_order(mock_get): - mock_get.return_value = None - - assert order_categories({"B", "A", "C"}, "en-US") == ["A", "B", "C"] diff --git a/api/tests/unit_tests/services/recommend_app/test_database_retrieval.py b/api/tests/unit_tests/services/recommend_app/test_database_retrieval.py deleted file mode 100644 index 44a6f782982..00000000000 --- a/api/tests/unit_tests/services/recommend_app/test_database_retrieval.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Unit tests for database recommendation retrieval delegation.""" - -from unittest.mock import patch - -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session - -from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval -from services.recommend_app.recommend_app_type import RecommendAppType - - -class TestDatabaseRecommendAppRetrieval: - def test_get_type(self) -> None: - assert DatabaseRecommendAppRetrieval().get_type() == RecommendAppType.DATABASE - - def test_get_recommended_apps_delegates(self, sqlite_engine: Engine) -> None: - with ( - Session(sqlite_engine) as session, - patch.object( - DatabaseRecommendAppRetrieval, - "fetch_recommended_apps_from_db", - return_value={"recommended_apps": [], "categories": []}, - ) as mock_fetch, - ): - result = DatabaseRecommendAppRetrieval().get_recommended_apps_and_categories("en-US", session=session) - - mock_fetch.assert_called_once_with("en-US", session=session) - assert result == {"recommended_apps": [], "categories": []} - - def test_get_recommend_app_detail_delegates(self, sqlite_engine: Engine) -> None: - with ( - Session(sqlite_engine) as session, - patch.object( - DatabaseRecommendAppRetrieval, - "fetch_recommended_app_detail_from_db", - return_value={"id": "app-1"}, - ) as mock_fetch, - ): - result = DatabaseRecommendAppRetrieval().get_recommend_app_detail("app-1", session=session) - - mock_fetch.assert_called_once_with("app-1", session=session) - assert result == {"id": "app-1"} diff --git a/api/tests/unit_tests/services/recommend_app/test_recommend_app_factory.py b/api/tests/unit_tests/services/recommend_app/test_recommend_app_factory.py deleted file mode 100644 index 036cba0cc00..00000000000 --- a/api/tests/unit_tests/services/recommend_app/test_recommend_app_factory.py +++ /dev/null @@ -1,28 +0,0 @@ -import pytest - -from services.recommend_app.buildin.buildin_retrieval import BuildInRecommendAppRetrieval -from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval -from services.recommend_app.recommend_app_factory import RecommendAppRetrievalFactory -from services.recommend_app.remote.remote_retrieval import RemoteRecommendAppRetrieval - - -class TestRecommendAppRetrievalFactory: - @pytest.mark.parametrize( - ("mode", "expected_class"), - [ - ("remote", RemoteRecommendAppRetrieval), - ("builtin", BuildInRecommendAppRetrieval), - ("db", DatabaseRecommendAppRetrieval), - ], - ) - def test_factory_returns_correct_class(self, mode, expected_class): - result = RecommendAppRetrievalFactory.get_recommend_app_factory(mode) - assert result is expected_class - - def test_factory_raises_for_unknown_mode(self): - with pytest.raises(ValueError, match="invalid fetch recommended apps mode"): - RecommendAppRetrievalFactory.get_recommend_app_factory("invalid_mode") - - def test_get_buildin_recommend_app_retrieval(self): - result = RecommendAppRetrievalFactory.get_buildin_recommend_app_retrieval() - assert result is BuildInRecommendAppRetrieval diff --git a/api/tests/unit_tests/services/recommend_app/test_recommend_app_type.py b/api/tests/unit_tests/services/recommend_app/test_recommend_app_type.py deleted file mode 100644 index 08f72a6f774..00000000000 --- a/api/tests/unit_tests/services/recommend_app/test_recommend_app_type.py +++ /dev/null @@ -1,18 +0,0 @@ -from services.recommend_app.recommend_app_type import RecommendAppType - - -def test_enum_values(): - assert RecommendAppType.REMOTE == "remote" - assert RecommendAppType.BUILDIN == "builtin" - assert RecommendAppType.DATABASE == "db" - - -def test_enum_membership(): - assert "remote" in RecommendAppType.__members__.values() - assert "builtin" in RecommendAppType.__members__.values() - assert "db" in RecommendAppType.__members__.values() - - -def test_enum_is_str(): - for member in RecommendAppType: - assert isinstance(member, str) diff --git a/api/tests/unit_tests/services/recommend_app/test_remote_retrieval.py b/api/tests/unit_tests/services/recommend_app/test_remote_retrieval.py deleted file mode 100644 index 165381c1f68..00000000000 --- a/api/tests/unit_tests/services/recommend_app/test_remote_retrieval.py +++ /dev/null @@ -1,296 +0,0 @@ -from collections.abc import Iterator -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask -from sqlalchemy import Engine -from sqlalchemy.orm import Session - -from services.recommend_app.recommend_app_type import RecommendAppType -from services.recommend_app.remote.remote_retrieval import RemoteRecommendAppRetrieval, clear_remote_fetch_cache - - -@pytest.fixture(autouse=True) -def _clear_remote_fetch_cache_between_tests(): - clear_remote_fetch_cache() - yield - clear_remote_fetch_cache() - - -@pytest.fixture -def empty_sqlite_session(sqlite_engine: Engine) -> Iterator[Session]: - with Session(sqlite_engine) as session: - yield session - - -class TestRemoteRecommendAppRetrieval: - def test_get_type(self): - assert RemoteRecommendAppRetrieval().get_type() == RecommendAppType.REMOTE - - @patch.object( - RemoteRecommendAppRetrieval, - "fetch_recommended_app_detail_from_dify_official", - return_value={"id": "app-1"}, - ) - def test_get_recommend_app_detail_success(self, mock_fetch, empty_sqlite_session: Session): - result = RemoteRecommendAppRetrieval().get_recommend_app_detail("app-1", session=empty_sqlite_session) - assert result == {"id": "app-1"} - mock_fetch.assert_called_once_with("app-1") - assert not empty_sqlite_session.in_transaction() - - @patch( - "services.recommend_app.remote.remote_retrieval" - ".BuildInRecommendAppRetrieval.fetch_recommended_app_detail_from_builtin", - return_value={"id": "fallback"}, - ) - @patch.object( - RemoteRecommendAppRetrieval, - "fetch_recommended_app_detail_from_dify_official", - side_effect=ConnectionError("timeout"), - ) - def test_get_recommend_app_detail_falls_back_on_error( - self, mock_fetch, mock_builtin, empty_sqlite_session: Session - ): - result = RemoteRecommendAppRetrieval().get_recommend_app_detail("app-1", session=empty_sqlite_session) - assert result == {"id": "fallback"} - mock_builtin.assert_called_once_with("app-1") - assert not empty_sqlite_session.in_transaction() - - @patch.object( - RemoteRecommendAppRetrieval, - "fetch_recommended_apps_from_dify_official", - return_value={"recommended_apps": [], "categories": []}, - ) - def test_get_recommended_apps_success(self, mock_fetch, empty_sqlite_session: Session): - result = RemoteRecommendAppRetrieval().get_recommended_apps_and_categories( - "en-US", session=empty_sqlite_session - ) - assert result == {"recommended_apps": [], "categories": []} - assert not empty_sqlite_session.in_transaction() - - @patch( - "services.recommend_app.remote.remote_retrieval" - ".BuildInRecommendAppRetrieval.fetch_recommended_apps_from_builtin", - return_value={"recommended_apps": [{"id": "builtin"}]}, - ) - @patch.object( - RemoteRecommendAppRetrieval, - "fetch_recommended_apps_from_dify_official", - side_effect=ValueError("server error"), - ) - def test_get_recommended_apps_falls_back_on_error(self, mock_fetch, mock_builtin, empty_sqlite_session: Session): - result = RemoteRecommendAppRetrieval().get_recommended_apps_and_categories( - "en-US", session=empty_sqlite_session - ) - assert result == {"recommended_apps": [{"id": "builtin"}]} - assert not empty_sqlite_session.in_transaction() - - @patch.object( - RemoteRecommendAppRetrieval, - "fetch_learn_dify_apps_from_dify_official", - return_value={"recommended_apps": [{"id": "learn-dify-app"}]}, - ) - def test_get_learn_dify_apps_success(self, mock_fetch, empty_sqlite_session: Session): - result = RemoteRecommendAppRetrieval().get_learn_dify_apps("en-US", session=empty_sqlite_session) - - assert result == {"recommended_apps": [{"id": "learn-dify-app"}]} - mock_fetch.assert_called_once_with("en-US") - assert not empty_sqlite_session.in_transaction() - - @patch( - "services.recommend_app.remote.remote_retrieval.DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db", - return_value={"recommended_apps": [{"id": "db-fallback"}]}, - ) - @patch.object( - RemoteRecommendAppRetrieval, - "fetch_learn_dify_apps_from_dify_official", - side_effect=ValueError("server error"), - ) - def test_get_learn_dify_apps_falls_back_to_database_on_error( - self, mock_fetch, mock_database, empty_sqlite_session: Session - ): - result = RemoteRecommendAppRetrieval().get_learn_dify_apps("en-US", session=empty_sqlite_session) - - assert result == {"recommended_apps": [{"id": "db-fallback"}]} - mock_database.assert_called_once_with("en-US", session=empty_sqlite_session) - assert not empty_sqlite_session.in_transaction() - - -class TestFetchFromDifyOfficial: - @pytest.fixture(autouse=True) - def _remote_config(self, config_overrides): - config_overrides( - HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN="https://example.com", - HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL=300, - CONSOLE_WEB_URL="", - ) - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_detail_returns_json_on_200(self, mock_get): - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"id": "app-1", "name": "Test"} - mock_get.return_value = mock_response - - result = RemoteRecommendAppRetrieval.fetch_recommended_app_detail_from_dify_official("app-1") - - assert result == {"id": "app-1", "name": "Test"} - mock_get.assert_called_once() - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_detail_returns_none_on_non_200(self, mock_get): - mock_get.return_value = MagicMock(status_code=404) - - result = RemoteRecommendAppRetrieval.fetch_recommended_app_detail_from_dify_official("app-1") - - assert result is None - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_preserves_remote_categories_order_on_200(self, mock_get): - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = { - "recommended_apps": [], - "categories": ["writing", "agent", "chat"], - } - mock_get.return_value = mock_response - - result = RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert result["categories"] == ["writing", "agent", "chat"] - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_raises_on_non_200(self, mock_get): - mock_get.return_value = MagicMock(status_code=500) - - with pytest.raises(ValueError, match="fetch recommended apps failed"): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_without_categories_key(self, mock_get): - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": []} - mock_get.return_value = mock_response - - result = RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert "categories" not in result - assert mock_get.call_args.kwargs["headers"] == {} - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_forwards_request_origin_header(self, mock_get, config_overrides): - config_overrides(CONSOLE_WEB_URL="https://saas.dify.dev") - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": []} - mock_get.return_value = mock_response - - flask_app = Flask(__name__) - with flask_app.test_request_context(headers={"Origin": "https://cloud.example.com"}): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert mock_get.call_args.kwargs["headers"] == {"Origin": "https://cloud.example.com"} - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_falls_back_to_console_web_url_origin(self, mock_get, config_overrides): - config_overrides(CONSOLE_WEB_URL="https://saas.dify.dev/console") - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": []} - mock_get.return_value = mock_response - - flask_app = Flask(__name__) - with flask_app.test_request_context(): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert mock_get.call_args.kwargs["headers"] == {"Origin": "https://saas.dify.dev/console"} - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_falls_back_to_console_web_url_without_request_context(self, mock_get, config_overrides): - config_overrides(CONSOLE_WEB_URL="http://localhost:3000/console") - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": []} - mock_get.return_value = mock_response - - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert mock_get.call_args.kwargs["headers"] == {"Origin": "http://localhost:3000/console"} - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_uses_console_web_url_without_scheme(self, mock_get, config_overrides): - config_overrides(CONSOLE_WEB_URL="saas.dify.dev") - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": []} - mock_get.return_value = mock_response - - flask_app = Flask(__name__) - with flask_app.test_request_context(): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert mock_get.call_args.kwargs["headers"] == {"Origin": "saas.dify.dev"} - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_learn_dify_apps_returns_json_on_200(self, mock_get): - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": [{"id": "learn-dify-app"}]} - mock_get.return_value = mock_response - - result = RemoteRecommendAppRetrieval.fetch_learn_dify_apps_from_dify_official("en-US") - - assert result == {"recommended_apps": [{"id": "learn-dify-app"}]} - assert mock_get.call_args.args[0] == "https://example.com/apps/learn-dify?language=en-US" - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_learn_dify_apps_raises_on_non_200(self, mock_get): - mock_get.return_value = MagicMock(status_code=500) - - with pytest.raises(ValueError, match="fetch learn dify apps failed"): - RemoteRecommendAppRetrieval.fetch_learn_dify_apps_from_dify_official("en-US") - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_uses_cache_for_repeated_requests(self, mock_get, config_overrides): - config_overrides(HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL=600) - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": [{"id": "app-1"}]} - mock_get.return_value = mock_response - - first = RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - second = RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert first == second == {"recommended_apps": [{"id": "app-1"}]} - mock_get.assert_called_once() - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_does_not_cache_failed_responses(self, mock_get, config_overrides): - config_overrides(HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL=600) - mock_get.return_value = MagicMock(status_code=500) - - with pytest.raises(ValueError, match="fetch recommended apps failed"): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - with pytest.raises(ValueError, match="fetch recommended apps failed"): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert mock_get.call_count == 2 - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_skips_cache_when_ttl_disabled(self, mock_get, config_overrides): - config_overrides(HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL=0) - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": []} - mock_get.return_value = mock_response - - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert mock_get.call_count == 2 - - @patch("services.recommend_app.remote.remote_retrieval.httpx.get") - def test_apps_cache_isolated_by_origin_header(self, mock_get, config_overrides): - config_overrides(HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL=600) - mock_response = MagicMock(status_code=200) - mock_response.json.return_value = {"recommended_apps": []} - mock_get.return_value = mock_response - - flask_app = Flask(__name__) - with flask_app.test_request_context(headers={"Origin": "https://cloud-a.example.com"}): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - with flask_app.test_request_context(headers={"Origin": "https://cloud-b.example.com"}): - RemoteRecommendAppRetrieval.fetch_recommended_apps_from_dify_official("en-US") - - assert mock_get.call_count == 2 diff --git a/api/tests/unit_tests/services/test_account_activation_adapters.py b/api/tests/unit_tests/services/test_account_activation_adapters.py index 931001197c7..402ee708d33 100644 --- a/api/tests/unit_tests/services/test_account_activation_adapters.py +++ b/api/tests/unit_tests/services/test_account_activation_adapters.py @@ -42,11 +42,22 @@ def test_invitation_token_store_revokes_with_legacy_key_inputs() -> None: def test_billing_eligibility_skips_gateway_when_disabled() -> None: - with patch("services.account_activation_adapters.BillingService.is_email_in_freeze") as is_frozen: - result = BillingAccountActivationEligibility(enabled=False).is_frozen("invitee@example.com") + with patch("services.account_activation_adapters.BillingService.get_email_freeze_type") as get_freeze_type: + result = BillingAccountActivationEligibility(enabled=False).get_freeze_type("invitee@example.com") - assert result is False - is_frozen.assert_not_called() + assert result is None + get_freeze_type.assert_not_called() + + +def test_billing_eligibility_returns_freeze_type_when_enabled() -> None: + with patch( + "services.account_activation_adapters.BillingService.get_email_freeze_type", + return_value="email_domain_suspended", + ) as get_freeze_type: + result = BillingAccountActivationEligibility(enabled=True).get_freeze_type("invitee@example.com") + + assert result == "email_domain_suspended" + get_freeze_type.assert_called_once_with("invitee@example.com") def test_membership_cache_skips_gateway_when_disabled() -> None: diff --git a/api/tests/unit_tests/services/test_account_activation_service.py b/api/tests/unit_tests/services/test_account_activation_service.py index e56214b3eb5..cd43170fb9e 100644 --- a/api/tests/unit_tests/services/test_account_activation_service.py +++ b/api/tests/unit_tests/services/test_account_activation_service.py @@ -6,6 +6,7 @@ from services.account_activation_service import ( AccountActivationEligibility, AccountActivationRepository, AccountActivationService, + EmailDomainSuspendedError, FrozenAccountError, InvalidInvitationError, InvitationAccountMismatchError, @@ -60,7 +61,7 @@ def _service() -> tuple[AccountActivationService, Mock, Mock, Mock, Mock, Mock]: policy = Mock(spec=WorkspaceInvitePolicy) eligibility = Mock(spec=AccountActivationEligibility) membership_cache = Mock(spec=WorkspaceMembershipCache) - eligibility.is_frozen.return_value = False + eligibility.get_freeze_type.return_value = None service = AccountActivationService( tokens=tokens, accounts=accounts, @@ -133,7 +134,7 @@ class TestActivateInvitation: authenticated_account_id="different-account", ) - eligibility.is_frozen.assert_not_called() + eligibility.get_freeze_type.assert_not_called() tokens.revoke.assert_not_called() accounts.activate.assert_not_called() @@ -141,12 +142,12 @@ class TestActivateInvitation: service, tokens, accounts, _, eligibility, _ = _service() tokens.find.return_value = _token() accounts.resolve.return_value = _invitation() - eligibility.is_frozen.return_value = True + eligibility.get_freeze_type.return_value = "freeze" with pytest.raises(FrozenAccountError): service.activate(ActivationCommand(invitation=_lookup()), authenticated_account_id=None) - eligibility.is_frozen.assert_called_once_with("invitee@example.com") + eligibility.get_freeze_type.assert_called_once_with("invitee@example.com") tokens.revoke.assert_not_called() accounts.activate.assert_not_called() @@ -164,6 +165,19 @@ class TestActivateInvitation: tokens.revoke.assert_not_called() accounts.activate.assert_not_called() + def test_rejects_suspended_email_domain_without_consuming_token(self) -> None: + service, tokens, accounts, _, eligibility, _ = _service() + tokens.find.return_value = _token() + accounts.resolve.return_value = _invitation() + eligibility.get_freeze_type.return_value = "email_domain_suspended" + + with pytest.raises(EmailDomainSuspendedError): + service.activate(ActivationCommand(invitation=_lookup()), authenticated_account_id=None) + + eligibility.get_freeze_type.assert_called_once_with("invitee@example.com") + tokens.revoke.assert_not_called() + accounts.activate.assert_not_called() + def test_activates_anonymous_invitation_and_invalidates_new_membership_cache(self) -> None: service, tokens, accounts, _, eligibility, membership_cache = _service() tokens.find.return_value = _token() @@ -179,7 +193,7 @@ class TestActivateInvitation: service.activate(command, authenticated_account_id=None) - eligibility.is_frozen.assert_called_once_with("invitee@example.com") + eligibility.get_freeze_type.assert_called_once_with("invitee@example.com") tokens.revoke.assert_called_once_with(_lookup("invitee@example.com")) accounts.activate.assert_called_once_with( invitation, diff --git a/api/tests/unit_tests/services/test_account_profile_service.py b/api/tests/unit_tests/services/test_account_profile_service.py new file mode 100644 index 00000000000..018cd45fc59 --- /dev/null +++ b/api/tests/unit_tests/services/test_account_profile_service.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from datetime import datetime +from unittest.mock import Mock + +import pytest + +from machinery.context import RequestContext +from services.account_errors import AccountNotFoundError +from services.account_ports import AccountRepository +from services.account_profile_service import AccountProfileService +from services.entities.account_entities import AccountProfileChanges, AccountSnapshot + + +def _context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id="workspace-1", + ) + + +def _account() -> AccountSnapshot: + return AccountSnapshot( + id="account-1", + name="Account", + email="account@example.com", + avatar=None, + is_password_set=False, + interface_language="en-US", + interface_theme="light", + timezone="UTC", + last_login_at=None, + last_login_ip=None, + status="active", + initialized_at=None, + created_at=datetime(2026, 1, 1), + ) + + +def test_get_returns_framework_neutral_account_snapshot() -> None: + accounts = Mock(spec=AccountRepository) + accounts.get.return_value = _account() + service = AccountProfileService(accounts=accounts) + + result = service.get(_context()) + + assert result == _account() + accounts.get.assert_called_once_with("account-1") + + +def test_update_applies_profile_changes() -> None: + accounts = Mock(spec=AccountRepository) + accounts.update_profile.return_value = _account() + service = AccountProfileService(accounts=accounts) + changes = AccountProfileChanges(name="Updated", timezone="Asia/Singapore") + + result = service.update(_context(), changes) + + assert result == _account() + accounts.update_profile.assert_called_once_with("account-1", changes) + + +def test_update_treats_empty_changes_as_noop() -> None: + accounts = Mock(spec=AccountRepository) + accounts.get.return_value = _account() + service = AccountProfileService(accounts=accounts) + + result = service.update(_context(), AccountProfileChanges()) + + assert result == _account() + accounts.get.assert_called_once_with("account-1") + accounts.update_profile.assert_not_called() + + +def test_update_rejects_missing_account() -> None: + accounts = Mock(spec=AccountRepository) + accounts.update_profile.return_value = None + service = AccountProfileService(accounts=accounts) + changes = AccountProfileChanges(name="Updated") + + with pytest.raises(AccountNotFoundError): + service.update(_context(), changes) + + accounts.update_profile.assert_called_once_with("account-1", changes) diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 7b0b4e9187d..5257a8b4e40 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -29,6 +29,7 @@ from services.errors.account import ( AccountPasswordError, AccountRegisterError, CurrentPasswordIncorrectError, + EmailDomainSuspendedError, NoPermissionError, ) @@ -341,6 +342,50 @@ class TestAccountService: session=unbound_session, ) + def test_create_account_suspended_email_domain( + self, unbound_session: Session, mock_external_service_dependencies: _MockDependencies + ) -> None: + mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True + mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = True + mock_external_service_dependencies[ + "billing_service" + ].get_email_freeze_type.return_value = "email_domain_suspended" + + with patch("services.account_service.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD): + with pytest.raises(EmailDomainSuspendedError): + AccountService.create_account( + email="user@suspended.example", + name="Test User", + interface_language="en-US", + session=unbound_session, + ) + + def test_get_user_through_email_rejects_suspended_email_domain( + self, unbound_session: Session, mock_external_service_dependencies: _MockDependencies + ) -> None: + mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = True + mock_external_service_dependencies[ + "billing_service" + ].get_email_freeze_type.return_value = "email_domain_suspended" + + with patch("services.account_service.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD): + with pytest.raises(EmailDomainSuspendedError): + AccountService.get_user_through_email("user@suspended.example", session=unbound_session) + + def test_get_account_freeze_type_is_enabled_only_for_cloud( + self, mock_external_service_dependencies: _MockDependencies + ) -> None: + mock_external_service_dependencies["billing_service"].get_email_freeze_type.return_value = "freeze" + + with patch("services.account_service.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD): + assert AccountService.get_account_freeze_type("frozen@example.com") == "freeze" + with patch("services.account_service.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY): + assert AccountService.get_account_freeze_type("frozen@example.com") is None + + mock_external_service_dependencies["billing_service"].get_email_freeze_type.assert_called_once_with( + "frozen@example.com" + ) + def test_create_account_without_password( self, sqlite_session_factory: sessionmaker[Session], diff --git a/api/tests/unit_tests/services/test_api_token_service.py b/api/tests/unit_tests/services/test_api_token_service.py index 3f7ae699088..24a187551a3 100644 --- a/api/tests/unit_tests/services/test_api_token_service.py +++ b/api/tests/unit_tests/services/test_api_token_service.py @@ -14,10 +14,29 @@ from werkzeug.exceptions import Unauthorized import services.api_token_service as api_token_service_module from models.engine import db +from models.enums import ApiTokenType from models.model import ApiToken from services.api_token_service import ApiTokenCache, CachedApiToken +def _api_token( + *, + token_id: str = "id-123", + app_id: str = "app-123", + tenant_id: str = "tenant-123", + token: str = "token-123", +) -> ApiToken: + """Create a mapped API token for cache and single-flight behavior tests.""" + return ApiToken( + id=token_id, + app_id=app_id, + tenant_id=tenant_id, + type=ApiTokenType.APP, + token=token, + last_used_at=None, + ) + + @pytest.fixture def api_token_db() -> Iterator[Session]: """Provide the production database extension with an isolated SQLite token table.""" @@ -168,7 +187,7 @@ class TestFetchTokenWithSingleFlight: def test_should_query_db_when_lock_acquired_and_cache_missed(self): auth_token = "token-123" scope = "app" - db_token = MagicMock() + db_token = _api_token() lock = MagicMock() lock.acquire.return_value = True @@ -187,7 +206,7 @@ class TestFetchTokenWithSingleFlight: def test_should_query_db_directly_when_lock_not_acquired(self): auth_token = "token-123" scope = "app" - db_token = MagicMock() + db_token = _api_token() lock = MagicMock() lock.acquire.return_value = False @@ -228,7 +247,7 @@ class TestFetchTokenWithSingleFlight: def test_should_fallback_to_db_query_when_lock_raises_exception(self): auth_token = "token-123" scope = "app" - db_token = MagicMock() + db_token = _api_token() lock = MagicMock() lock.acquire.side_effect = RuntimeError("redis lock error") @@ -366,14 +385,7 @@ class TestApiTokenCacheCoreBranches: @patch("services.api_token_service.redis_client") def test_set_should_return_false_when_cache_write_raises_exception(self, mock_redis): mock_redis.setex.side_effect = Exception("redis write failed") - api_token = MagicMock() - api_token.id = "id-123" - api_token.app_id = "app-123" - api_token.tenant_id = "tenant-123" - api_token.type = "app" - api_token.token = "token-123" - api_token.last_used_at = None - api_token.created_at = None + api_token = _api_token() result = ApiTokenCache.set("token-123", "app", api_token) assert result is False diff --git a/api/tests/unit_tests/services/test_app_dsl_service.py b/api/tests/unit_tests/services/test_app_dsl_service.py index 7afbb0fa042..9b76c0938c4 100644 --- a/api/tests/unit_tests/services/test_app_dsl_service.py +++ b/api/tests/unit_tests/services/test_app_dsl_service.py @@ -7,7 +7,7 @@ import yaml from sqlalchemy import event, select from sqlalchemy.orm import Session, sessionmaker -from core.rbac import RBACPermission +from core.rbac import RBACPermission, RBACResourceScope from core.workflow.llm_environment_variable import LLMEnvironmentVariable from models import App, AppMode from models.model import AppModelConfig, AppModelConfigDict, IconType @@ -17,6 +17,40 @@ from services.entities.dsl_entities import ImportStatus from services.errors.account import NoPermissionError from services.errors.app import WorkflowNotFoundError +_OVERWRITE_APP_ID = "11111111-1111-4111-8111-111111111111" +_TENANT_ID = "22222222-2222-4222-8222-222222222222" +_CALLER_ID = "33333333-3333-4333-8333-333333333333" +_OTHER_ACCOUNT_ID = "44444444-4444-4444-8444-444444444444" +_PENDING_WORKFLOW_DSL = "version: 99.0.0\nkind: app\napp: {name: Test, mode: workflow}\n" +_PENDING_DATA_JSON = PendingData( + tenant_id=_TENANT_ID, + account_id=_CALLER_ID, + import_mode="yaml-content", + yaml_content=_PENDING_WORKFLOW_DSL, + app_id=_OVERWRITE_APP_ID, +).model_dump_json() + + +def _persist_overwrite_target(session: Session, *, maintainer: str = _OTHER_ACCOUNT_ID) -> App: + app = App( + id=_OVERWRITE_APP_ID, + tenant_id=_TENANT_ID, + name="Target", + description="", + mode=AppMode.WORKFLOW, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + created_by=maintainer, + maintainer=maintainer, + updated_by=maintainer, + ) + session.add(app) + session.commit() + return app + def test_extract_workflow_dependencies_uses_llm_environment_variable_provider(monkeypatch: pytest.MonkeyPatch) -> None: workflow = SimpleNamespace( @@ -145,6 +179,99 @@ def test_import_app_returns_decode_error_for_invalid_yaml_url_bytes( assert not unbound_session.in_transaction() +def test_import_app_checks_overwrite_rbac_before_database_access( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + _persist_overwrite_target(sqlite_session) + account = Mock(id=_CALLER_ID, current_tenant_id=_TENANT_ID) + + def deny_before_transaction(*_args: object, **_kwargs: object) -> bool: + assert not sqlite_session.in_transaction() + return False + + check = Mock(side_effect=deny_before_transaction) + setex = Mock() + monkeypatch.setattr("services.app_dsl_service.dify_config.RBAC_ENABLED", True) + monkeypatch.setattr("services.app_dsl_service.RBACService.CheckAccess.check", check) + monkeypatch.setattr("services.app_dsl_service.redis_client.setex", setex) + + with pytest.raises(NoPermissionError, match="permission to overwrite"): + AppDslService(sqlite_session).import_app( + account=account, + import_mode="yaml-content", + yaml_content=_PENDING_WORKFLOW_DSL, + app_id=_OVERWRITE_APP_ID, + ) + + check.assert_called_once_with( + _TENANT_ID, + _CALLER_ID, + scene=RBACPermission.APP_IMPORT_EXPORT_DSL, + resource_type=RBACResourceScope.APP, + resource_id=_OVERWRITE_APP_ID, + ) + setex.assert_not_called() + + +def test_confirm_import_rechecks_overwrite_rbac_before_database_access( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + _persist_overwrite_target(sqlite_session) + monkeypatch.setattr("services.app_dsl_service.redis_client.get", Mock(return_value=_PENDING_DATA_JSON)) + redis_delete = Mock() + monkeypatch.setattr("services.app_dsl_service.redis_client.delete", redis_delete) + create_or_update = Mock() + service = AppDslService(sqlite_session) + monkeypatch.setattr(service, "_create_or_update_app", create_or_update) + + def deny_before_transaction(*_args: object, **_kwargs: object) -> bool: + assert not sqlite_session.in_transaction() + return False + + check = Mock(side_effect=deny_before_transaction) + monkeypatch.setattr("services.app_dsl_service.dify_config.RBAC_ENABLED", True) + monkeypatch.setattr("services.app_dsl_service.RBACService.CheckAccess.check", check) + + with pytest.raises(NoPermissionError, match="permission to overwrite"): + service.confirm_import( + import_id="import-1", + account=Mock(id=_CALLER_ID, current_tenant_id=_TENANT_ID), + ) + + check.assert_called_once_with( + _TENANT_ID, + _CALLER_ID, + scene=RBACPermission.APP_IMPORT_EXPORT_DSL, + resource_type=RBACResourceScope.APP, + resource_id=_OVERWRITE_APP_ID, + ) + create_or_update.assert_not_called() + redis_delete.assert_not_called() + + +def test_confirm_import_does_not_create_when_overwrite_target_disappeared( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + monkeypatch.setattr("services.app_dsl_service.redis_client.get", Mock(return_value=_PENDING_DATA_JSON)) + monkeypatch.setattr("services.app_dsl_service.dify_config.RBAC_ENABLED", True) + monkeypatch.setattr("services.app_dsl_service.RBACService.CheckAccess.check", Mock(return_value=True)) + redis_delete = Mock() + monkeypatch.setattr("services.app_dsl_service.redis_client.delete", redis_delete) + service = AppDslService(sqlite_session) + create_or_update = Mock() + monkeypatch.setattr(service, "_create_or_update_app", create_or_update) + + result = service.confirm_import( + import_id="import-1", + account=Mock(id=_CALLER_ID, current_tenant_id=_TENANT_ID), + ) + + assert result.status == ImportStatus.FAILED + assert result.error == "App not found" + create_or_update.assert_not_called() + redis_delete.assert_not_called() + + def test_pending_import_is_scoped_to_its_owner(monkeypatch: pytest.MonkeyPatch, unbound_session: Session) -> None: pending_imports: dict[str, str] = {} monkeypatch.setattr( @@ -301,6 +428,59 @@ def test_create_or_update_app_flushes_new_model_config_before_signal( assert sqlite_session.in_transaction() +def test_create_or_update_app_forwards_imported_agent_purge_ids(monkeypatch: pytest.MonkeyPatch) -> None: + session = cast(Session, SimpleNamespace(add=Mock(), flush=Mock(), commit=Mock(), get=Mock())) + service = AppDslService(session=session) + app = SimpleNamespace( + id="app-1", + tenant_id="tenant-1", + name="Workflow", + description="", + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + ) + workflow = SimpleNamespace(id="workflow-1") + workflow_service = SimpleNamespace( + get_draft_workflow=Mock(return_value=None), + sync_draft_workflow=Mock(return_value=workflow), + ) + monkeypatch.setattr("services.app_dsl_service.WorkflowService", Mock(return_value=workflow_service)) + monkeypatch.setattr( + "services.app_dsl_service.AgentDslService.graph_without_package_bindings", + Mock(return_value={"nodes": [], "edges": []}), + ) + monkeypatch.setattr( + "services.app_dsl_service.AgentDslService.import_workflow_packages", + Mock(return_value=(workflow, [], {"retired-agent"})), + ) + monkeypatch.setattr( + "services.app_dsl_service.WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync", + Mock(), + ) + retire_unowned = Mock() + monkeypatch.setattr( + "services.app_dsl_service.WorkflowAgentRetirementService.retire_unowned", + retire_unowned, + ) + + service._create_or_update_app( + app=cast(App, app), + data={ + "app": {"mode": AppMode.WORKFLOW.value}, + "workflow": {"graph": {"nodes": [], "edges": []}}, + "agent_packages": {"package-1": {}}, + }, + account=Mock(id="account-1"), + ) + + retire_unowned.assert_called_once_with( + tenant_id="tenant-1", + agent_ids={"retired-agent"}, + account_id="account-1", + ) + + def test_export_dsl_loads_model_config_and_annotation_reply_with_request_session( monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session], diff --git a/api/tests/unit_tests/services/test_app_service.py b/api/tests/unit_tests/services/test_app_service.py index f33c3082e52..223be9f678b 100644 --- a/api/tests/unit_tests/services/test_app_service.py +++ b/api/tests/unit_tests/services/test_app_service.py @@ -14,8 +14,17 @@ from enums import DeploymentEdition from graphon.model_runtime.entities.model_entities import ModelType from models import Account, Tenant from models.account import TenantAccountJoin, TenantAccountRole -from models.agent import Agent, AgentIconType, AgentScope, AgentSource, AgentStatus +from models.agent import ( + Agent, + AgentIconType, + AgentScope, + AgentSource, + AgentStatus, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) from models.model import App, AppMode, AppModelConfig, IconType +from models.workflow import Workflow, WorkflowType from services.agent.errors import AgentAccessNotReadyError, AgentNameConflictError from services.app_service import AppListParams, AppService, CreateAppParams @@ -615,6 +624,31 @@ class TestAgentAppType: def test_delete_agent_app_archives_backing_agent(self, sqlite_session: Session): app, backing_agent = _persist_agent_app(sqlite_session) + workflow_app = _persist_app(sqlite_session, tenant_id=app.tenant_id, name="Workflow") + workflow_app.mode = AppMode.WORKFLOW + referencing_workflow = Workflow.new( + tenant_id=app.tenant_id, + app_id=workflow_app.id, + type=WorkflowType.WORKFLOW.value, + version=Workflow.VERSION_DRAFT, + graph="{}", + features="{}", + created_by="account-1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow_binding = WorkflowAgentNodeBinding( + tenant_id=app.tenant_id, + app_id=workflow_app.id, + workflow_id=referencing_workflow.id, + workflow_version=referencing_workflow.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id=backing_agent.id, + current_snapshot_id="snapshot-1", + node_job_config={}, + ) workflow_agents = [ Agent( tenant_id=app.tenant_id, @@ -630,7 +664,7 @@ class TestAgentAppType: ) for index in range(2) ] - sqlite_session.add_all(workflow_agents) + sqlite_session.add_all([referencing_workflow, workflow_binding, *workflow_agents]) sqlite_session.commit() account_id = str(uuid4()) events: list[str] = [] @@ -642,7 +676,10 @@ class TestAgentAppType: patch("services.app_service.BillingService"), patch("services.app_service.EnterpriseService"), patch("services.app_service.FeatureService"), - patch("services.app_service.remove_app_and_related_data_task"), + patch( + "services.app_service.remove_app_and_related_data_task.delay", + side_effect=lambda **_kwargs: events.append("enqueue-app-cleanup"), + ), patch( "services.app_service.AgentHomeSnapshotService.retire_all_for_agent", return_value=["home-1"], @@ -653,9 +690,7 @@ class TestAgentAppType: ) as mock_retire_workspaces, patch( "services.app_service.WorkflowAgentRetirementService.retire_unowned", - side_effect=lambda **_kwargs: ( - events.append("retire-workflow-agents") or (["workflow-binding-1"], ["workflow-home-1"]) - ), + side_effect=lambda **_kwargs: events.append("retire-workflow-agents"), ) as mock_workflow_retirement, patch( "services.app_service.enqueue_agent_resource_collection", @@ -664,7 +699,13 @@ class TestAgentAppType: ): AppService().delete_app(app, session=sqlite_session) - assert events == ["retire-app-workspaces", "commit", "retire-workflow-agents", "enqueue"] + assert events == [ + "retire-app-workspaces", + "commit", + "enqueue-app-cleanup", + "retire-workflow-agents", + "enqueue", + ] sqlite_session.expire_all() persisted_agent = sqlite_session.get(Agent, backing_agent.id) assert persisted_agent is not None @@ -672,9 +713,12 @@ class TestAgentAppType: assert persisted_agent.status == AgentStatus.ARCHIVED assert persisted_agent.archived_by == account_id assert persisted_agent.archived_at is not None + persisted_workflow_binding = sqlite_session.get(WorkflowAgentNodeBinding, workflow_binding.id) + assert persisted_workflow_binding is not None + assert persisted_workflow_binding.agent_id == backing_agent.id mock_workflow_retirement.assert_called_once_with( tenant_id=app.tenant_id, - agent_ids=[agent.id for agent in workflow_agents], + agent_ids={agent.id for agent in workflow_agents}, account_id=account_id, ) mock_retire_workspaces.assert_called_once_with( @@ -690,8 +734,9 @@ class TestAgentAppType: mock_enqueue_collection.assert_called_once_with( tenant_id=app.tenant_id, workspace_ids=["workspace-1"], - binding_ids=["workflow-binding-1"], - home_snapshot_ids=["home-1", "workflow-home-1"], + binding_ids=[], + home_snapshot_ids=["home-1"], + purge_agent_ids=[backing_agent.id], ) def test_delete_app_commit_failure_does_not_retire_workflow_agents_or_enqueue(self, sqlite_session: Session): @@ -722,9 +767,100 @@ class TestAgentAppType: patch("services.app_service.AgentWorkspaceService.retire_all_for_app", return_value=["workspace-1"]), patch("services.app_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, patch("services.app_service.enqueue_agent_resource_collection") as enqueue_collection, + patch("services.app_service.remove_app_and_related_data_task.delay") as enqueue_app_cleanup, ): with pytest.raises(RuntimeError, match="commit failed"): AppService().delete_app(app, session=sqlite_session) retire_unowned.assert_not_called() enqueue_collection.assert_not_called() + enqueue_app_cleanup.assert_not_called() + + def test_delete_workflow_app_releases_all_bindings_before_retirement(self, sqlite_session: Session): + app = _persist_app(sqlite_session, tenant_id=str(uuid4())) + app.mode = AppMode.WORKFLOW + workflow = Workflow.new( + tenant_id=app.tenant_id, + app_id=app.id, + type=WorkflowType.WORKFLOW.value, + version="historical-version", + graph="{}", + features="{}", + created_by="account-1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + inline_binding = WorkflowAgentNodeBinding( + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id="inline-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="inline-agent", + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + roster_binding = WorkflowAgentNodeBinding( + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id="roster-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="roster-agent", + current_snapshot_id="snapshot-2", + node_job_config={}, + ) + sqlite_session.add_all([workflow, inline_binding, roster_binding]) + sqlite_session.commit() + events: list[str] = [] + + def retire_unowned(**kwargs): + events.append("retire") + assert sqlite_session.get(WorkflowAgentNodeBinding, inline_binding.id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, roster_binding.id) is None + assert kwargs["agent_ids"] == {"inline-agent"} + + with ( + patch("services.app_service.current_user", _account_identity(str(uuid4()))), + patch("services.app_service.app_was_deleted.send"), + patch("services.app_service.FeatureService"), + patch("services.app_service.BillingService"), + patch("services.app_service.EnterpriseService"), + patch("services.app_service.AgentWorkspaceService.retire_all_for_app", return_value=[]), + patch( + "services.app_service.remove_app_and_related_data_task.delay", + side_effect=lambda **_kwargs: events.append("enqueue-app-cleanup"), + ), + patch( + "services.app_service.WorkflowAgentRetirementService.retire_unowned", + side_effect=retire_unowned, + ), + patch("services.app_service.enqueue_agent_resource_collection"), + ): + AppService().delete_app(app, session=sqlite_session) + + assert events == ["enqueue-app-cleanup", "retire"] + assert sqlite_session.get(WorkflowAgentNodeBinding, inline_binding.id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, roster_binding.id) is None + + def test_delete_app_cleanup_enqueue_failure_propagates_before_retirement(self, sqlite_session: Session): + app = _persist_app(sqlite_session, tenant_id=str(uuid4())) + error = RuntimeError("broker unavailable") + + with ( + patch("services.app_service.current_user", _account_identity(str(uuid4()))), + patch("services.app_service.app_was_deleted.send"), + patch("services.app_service.AgentWorkspaceService.retire_all_for_app", return_value=[]), + patch("services.app_service.remove_app_and_related_data_task.delay", side_effect=error), + patch("services.app_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + patch("services.app_service.enqueue_agent_resource_collection") as enqueue_collection, + ): + with pytest.raises(RuntimeError) as exc_info: + AppService().delete_app(app, session=sqlite_session) + + assert exc_info.value is error + retire_unowned.assert_not_called() + enqueue_collection.assert_not_called() diff --git a/api/tests/unit_tests/services/test_async_workflow_service.py b/api/tests/unit_tests/services/test_async_workflow_service.py index 93599363c4f..403ee39d5b8 100644 --- a/api/tests/unit_tests/services/test_async_workflow_service.py +++ b/api/tests/unit_tests/services/test_async_workflow_service.py @@ -6,7 +6,10 @@ from unittest.mock import MagicMock, patch import pytest import services.async_workflow_service as async_workflow_service_module -from models.enums import AppTriggerType, CreatorUserRole, WorkflowRunTriggeredFrom, WorkflowTriggerStatus +from models.account import Account +from models.enums import AppTriggerType, CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom, WorkflowTriggerStatus +from models.model import App, AppMode, EndUser +from models.workflow import Workflow, WorkflowType from services.async_workflow_service import AsyncWorkflowService from services.errors.app import QuotaExceededError, WorkflowNotFoundError from services.workflow.entities import AsyncTriggerResponse, TriggerData @@ -36,6 +39,20 @@ class AsyncWorkflowServiceTestDataFactory: trigger_metadata=None, ) + @staticmethod + def create_app(app_id: str = "app-123", tenant_id: str = "tenant-123") -> App: + """Create an app model for service-return tests.""" + return App( + id=app_id, + tenant_id=tenant_id, + name="Async workflow app", + description="", + mode=AppMode.WORKFLOW, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + @staticmethod def create_trigger_log_with_data(trigger_data: TriggerData, retry_count: int = 0) -> MagicMock: """Create a mock trigger log with serialized trigger data.""" @@ -48,6 +65,41 @@ class AsyncWorkflowServiceTestDataFactory: trigger_log.to_dict.return_value = {"id": trigger_log.id} return trigger_log + @staticmethod + def create_workflow( + *, workflow_id: str = "workflow-123", app_id: str = "app-123", tenant_id: str = "tenant-123" + ) -> Workflow: + """Create a mapped workflow for service-return and trigger-log tests.""" + return Workflow( + id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + type=WorkflowType.WORKFLOW, + version="1", + graph="{}", + _features="{}", + created_by="account-123", + ) + + @staticmethod + def create_account(account_id: str = "account-123") -> Account: + """Create a mapped account for role-discrimination tests.""" + account = Account(name="Account", email=f"{account_id}@example.com") + account.id = account_id + return account + + @staticmethod + def create_end_user(end_user_id: str = "end-user-123") -> EndUser: + """Create a mapped end user for role-discrimination and retry tests.""" + return EndUser( + id=end_user_id, + tenant_id="tenant-123", + app_id="app-123", + type=EndUserType.BROWSER, + name="End User", + session_id=f"session-{end_user_id}", + ) + class TestAsyncWorkflowService: @pytest.fixture @@ -129,8 +181,7 @@ class TestAsyncWorkflowService: app_model.id = "app-123" session.scalar.return_value = app_model trigger_data = AsyncWorkflowServiceTestDataFactory.create_trigger_data() - workflow = MagicMock() - workflow.id = "workflow-123" + workflow = AsyncWorkflowServiceTestDataFactory.create_workflow() mocks = async_workflow_trigger_mocks mocks["dispatcher"].get_queue_name.return_value = queue_name @@ -145,15 +196,10 @@ class TestAsyncWorkflowService: quota_charge_mock = MagicMock() mocks["quota_service"].reserve.return_value = quota_charge_mock - class DummyAccount: - def __init__(self, user_id: str): - self.id = user_id + user = AsyncWorkflowServiceTestDataFactory.create_account() - with patch.object(async_workflow_service_module, "Account", DummyAccount): - user = DummyAccount("account-123") - - # Act - result = AsyncWorkflowService.trigger_workflow_async(session=session, user=user, trigger_data=trigger_data) + # Act + result = AsyncWorkflowService.trigger_workflow_async(session=session, user=user, trigger_data=trigger_data) # Assert assert isinstance(result, AsyncTriggerResponse) @@ -195,8 +241,7 @@ class TestAsyncWorkflowService: app_model.id = "app-123" session.scalar.return_value = app_model trigger_data = AsyncWorkflowServiceTestDataFactory.create_trigger_data() - workflow = MagicMock() - workflow.id = "workflow-123" + workflow = AsyncWorkflowServiceTestDataFactory.create_workflow() mocks = async_workflow_trigger_mocks mocks["dispatcher"].get_queue_name.return_value = QueuePriority.SANDBOX @@ -205,7 +250,7 @@ class TestAsyncWorkflowService: task_result = MagicMock(id="task-123") mocks["sandbox_task"].delay.return_value = task_result - user = SimpleNamespace(id="end-user-123") + user = AsyncWorkflowServiceTestDataFactory.create_end_user() # Act AsyncWorkflowService.trigger_workflow_async(session=session, user=user, trigger_data=trigger_data) @@ -231,7 +276,7 @@ class TestAsyncWorkflowService: with pytest.raises(WorkflowNotFoundError, match="App not found: missing-app"): AsyncWorkflowService.trigger_workflow_async( session=session, - user=SimpleNamespace(id="user-123"), + user=AsyncWorkflowServiceTestDataFactory.create_end_user("user-123"), trigger_data=trigger_data, ) @@ -246,8 +291,7 @@ class TestAsyncWorkflowService: app_model.id = "app-123" session.scalar.return_value = app_model trigger_data = AsyncWorkflowServiceTestDataFactory.create_trigger_data() - workflow = MagicMock() - workflow.id = "workflow-123" + workflow = AsyncWorkflowServiceTestDataFactory.create_workflow() mocks = async_workflow_trigger_mocks mocks["dispatcher"].get_queue_name.return_value = QueuePriority.TEAM @@ -263,7 +307,7 @@ class TestAsyncWorkflowService: with pytest.raises(QuotaExceededError) as exc_info: AsyncWorkflowService.trigger_workflow_async( session=session, - user=SimpleNamespace(id="user-123"), + user=AsyncWorkflowServiceTestDataFactory.create_end_user("user-123"), trigger_data=trigger_data, ) @@ -294,7 +338,7 @@ class TestAsyncWorkflowService: with pytest.raises(ValueError, match="Trigger log not found: missing-log"): AsyncWorkflowService.reinvoke_trigger( session=session, - user=SimpleNamespace(id="user-123"), + user=AsyncWorkflowServiceTestDataFactory.create_end_user("user-123"), workflow_trigger_log_id="missing-log", ) @@ -322,7 +366,7 @@ class TestAsyncWorkflowService: return_value=expected_response, ) as mock_trigger_workflow_async, ): - user = SimpleNamespace(id="user-123") + user = AsyncWorkflowServiceTestDataFactory.create_end_user("user-123") # Act response = AsyncWorkflowService.reinvoke_trigger( @@ -468,8 +512,8 @@ class TestAsyncWorkflowServiceGetWorkflow: """Test _get_workflow returns published workflow by id when provided.""" # Arrange workflow_service = MagicMock() - app_model = MagicMock() - workflow = MagicMock() + app_model = AsyncWorkflowServiceTestDataFactory.create_app() + workflow = AsyncWorkflowServiceTestDataFactory.create_workflow() workflow_service.get_published_workflow_by_id.return_value = workflow # Act @@ -489,7 +533,7 @@ class TestAsyncWorkflowServiceGetWorkflow: """Test _get_workflow raises WorkflowNotFoundError for unknown workflow id.""" # Arrange workflow_service = MagicMock() - app_model = MagicMock() + app_model = AsyncWorkflowServiceTestDataFactory.create_app() workflow_service.get_published_workflow_by_id.return_value = None # Act / Assert @@ -502,9 +546,8 @@ class TestAsyncWorkflowServiceGetWorkflow: """Test _get_workflow returns default published workflow when no id is provided.""" # Arrange workflow_service = MagicMock() - app_model = MagicMock() - app_model.id = "app-123" - workflow = MagicMock() + app_model = AsyncWorkflowServiceTestDataFactory.create_app() + workflow = AsyncWorkflowServiceTestDataFactory.create_workflow() workflow_service.get_published_workflow.return_value = workflow # Act @@ -520,8 +563,7 @@ class TestAsyncWorkflowServiceGetWorkflow: """Test _get_workflow raises WorkflowNotFoundError when app has no published workflow.""" # Arrange workflow_service = MagicMock() - app_model = MagicMock() - app_model.id = "app-123" + app_model = AsyncWorkflowServiceTestDataFactory.create_app() workflow_service.get_published_workflow.return_value = None # Act / Assert diff --git a/api/tests/unit_tests/services/test_audio_service.py b/api/tests/unit_tests/services/test_audio_service.py index e4d6046a3ea..33612604fbe 100644 --- a/api/tests/unit_tests/services/test_audio_service.py +++ b/api/tests/unit_tests/services/test_audio_service.py @@ -53,9 +53,11 @@ Tests available voice retrieval: - text_to_speech: Enables TTS functionality """ +import json from decimal import Decimal from typing import Any -from unittest.mock import MagicMock, Mock, create_autospec, patch +from unittest.mock import MagicMock, Mock, patch +from uuid import uuid4 import pytest from sqlalchemy.orm import Session @@ -64,7 +66,7 @@ from werkzeug.datastructures import FileStorage from models.agent_config_entities import AgentSoulConfig from models.enums import ConversationFromSource, MessageStatus from models.model import App, AppMode, AppModelConfig, Message -from models.workflow import Workflow +from models.workflow import Workflow, WorkflowType from services.app_ref_service import AppRef, MessageRef from services.audio_service import AudioService from services.errors.audio import ( @@ -113,15 +115,18 @@ class AudioServiceTestDataFactory: audio-related operations. """ - @staticmethod + def __init__(self, session: Session) -> None: + self.session = session + def create_app_mock( - app_id: str = "app-123", + self, + app_id: str = APP_ID, mode: AppMode = AppMode.CHAT, - tenant_id: str = "tenant-123", + tenant_id: str = TENANT_ID, **kwargs, - ) -> Mock: + ) -> App: """ - Create a mock App object. + Create and persist an App model. Args: app_id: Unique identifier for the app @@ -130,46 +135,65 @@ class AudioServiceTestDataFactory: **kwargs: Additional attributes to set on the mock Returns: - Mock App object with specified attributes + Persisted App model with specified attributes """ - app = create_autospec(App, instance=True) - app.id = app_id - app.mode = mode - app.tenant_id = tenant_id - app.workflow = kwargs.get("workflow") - app.app_model_config = kwargs.get("app_model_config") - app.workflow_with_session.return_value = app.workflow - app.app_model_config_with_session.return_value = app.app_model_config + workflow = kwargs.pop("workflow", None) + app_model_config = kwargs.pop("app_model_config", None) + app = App( + id=app_id, + tenant_id=tenant_id, + name="Audio test app", + description="", + mode=mode, + icon_type=None, + icon=None, + icon_background=None, + enable_site=False, + enable_api=False, + workflow_id=workflow.id if workflow else None, + app_model_config_id=app_model_config.id if app_model_config else None, + ) for key, value in kwargs.items(): setattr(app, key, value) + self.session.add(app) + self.session.commit() return app - @staticmethod - def create_workflow_mock(features_dict: dict[str, Any] | None = None, **kwargs) -> Mock: + def create_workflow_mock(self, features_dict: dict[str, Any] | None = None, **kwargs) -> Workflow: """ - Create a mock Workflow object. + Create and persist a Workflow model. Args: features_dict: Dictionary of workflow features **kwargs: Additional attributes to set on the mock Returns: - Mock Workflow object with specified attributes + Persisted Workflow model with specified attributes """ - workflow = create_autospec(Workflow, instance=True) - workflow.features_dict = features_dict or {} + workflow = Workflow( + id=kwargs.pop("id", str(uuid4())), + tenant_id=kwargs.pop("tenant_id", TENANT_ID), + app_id=kwargs.pop("app_id", APP_ID), + type=kwargs.pop("type", WorkflowType.CHAT), + version=kwargs.pop("version", Workflow.VERSION_DRAFT), + graph=kwargs.pop("graph", "{}"), + _features=json.dumps(features_dict or {}), + created_by=kwargs.pop("created_by", ACCOUNT_ID), + ) for key, value in kwargs.items(): setattr(workflow, key, value) + self.session.add(workflow) + self.session.commit() return workflow - @staticmethod def create_app_model_config_mock( + self, speech_to_text_dict: dict[str, Any] | None = None, text_to_speech_dict: dict[str, Any] | None = None, **kwargs, - ) -> Mock: + ) -> AppModelConfig: """ - Create a mock AppModelConfig object. + Create and persist an AppModelConfig model. Args: speech_to_text_dict: Speech-to-text configuration @@ -177,13 +201,17 @@ class AudioServiceTestDataFactory: **kwargs: Additional attributes to set on the mock Returns: - Mock AppModelConfig object with specified attributes + Persisted AppModelConfig model with specified attributes """ - config = create_autospec(AppModelConfig, instance=True) - config.speech_to_text_dict = speech_to_text_dict or {"enabled": False} - config.text_to_speech_dict = text_to_speech_dict or {"enabled": False} + config = AppModelConfig( + app_id=kwargs.pop("app_id", APP_ID), + speech_to_text=json.dumps(speech_to_text_dict or {"enabled": False}), + text_to_speech=json.dumps(text_to_speech_dict or {"enabled": False}), + ) for key, value in kwargs.items(): setattr(config, key, value) + self.session.add(config) + self.session.commit() return config @staticmethod @@ -216,9 +244,9 @@ class AudioServiceTestDataFactory: @pytest.fixture -def factory(): +def factory(sqlite_session: Session) -> AudioServiceTestDataFactory: """Provide the test data factory to all tests.""" - return AudioServiceTestDataFactory + return AudioServiceTestDataFactory(sqlite_session) class TestAudioServiceASR: @@ -385,7 +413,6 @@ class TestAudioServiceASR: self, mock_model_manager_class, factory: AudioServiceTestDataFactory ): app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app_model_config.to_dict.return_value = {"speech_to_text": {"enabled": True}} app = factory.create_app_mock(mode=AppMode.AGENT, app_model_config=app_model_config) file = factory.create_file_storage_mock() mock_model_instance = MagicMock() @@ -403,7 +430,6 @@ class TestAudioServiceASR: def test_transcript_agent_asr_soul_disabled_overrides_legacy_feature(self, factory: AudioServiceTestDataFactory): app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app_model_config.to_dict.return_value = {"speech_to_text": {"enabled": True}} app = factory.create_app_mock(mode=AppMode.AGENT, app_model_config=app_model_config) file = factory.create_file_storage_mock() agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": False}}}) diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index 5b1ce38ad4a..22a50d3be6b 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -1297,6 +1297,15 @@ class TestBillingServiceAccountManagement: assert result is True mock_send_request.assert_called_once_with("GET", "/account/in-freeze", params={"email": email}) + def test_get_email_freeze_type_for_suspended_domain(self, mock_send_request): + email = "user@suspended.example" + mock_send_request.return_value = {"data": True, "freezeType": "email_domain_suspended"} + + result = BillingService.get_email_freeze_type(email) + + assert result == "email_domain_suspended" + mock_send_request.assert_called_once_with("GET", "/account/in-freeze", params={"email": email}) + def test_is_email_in_freeze_false(self, mock_send_request): """Test checking if email is frozen (returns False).""" # Arrange diff --git a/api/tests/unit_tests/services/test_dataset_service_document.py b/api/tests/unit_tests/services/test_dataset_service_document.py index 7093053d1cb..1763cfab7b9 100644 --- a/api/tests/unit_tests/services/test_dataset_service_document.py +++ b/api/tests/unit_tests/services/test_dataset_service_document.py @@ -1,5 +1,15 @@ """Unit tests for DocumentService behaviors in dataset_service.""" +from datetime import datetime + +from sqlalchemy import event, select +from sqlalchemy.orm import Session + +from models.account import Tenant +from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from models.model import UploadFile +from models.source import DataSourceOauthBinding from services.dataset_ref_service import DatasetRefService from .dataset_service_test_helpers import ( @@ -8,7 +18,6 @@ from .dataset_service_test_helpers import ( CloudPlan, DatasetProcessRule, DatasetService, - DatasetServiceUnitDataFactory, DataSource, Document, DocumentIndexingError, @@ -18,7 +27,6 @@ from .dataset_service_test_helpers import ( IndexStructureType, InfoList, KnowledgeConfig, - MagicMock, NotFound, NotionIcon, NotionInfo, @@ -32,18 +40,117 @@ from .dataset_service_test_helpers import ( Segmentation, SimpleNamespace, WebsiteInfo, - _make_dataset, - _make_document, _make_features, _make_lock_context, _make_upload_knowledge_config, - create_autospec, json, patch, pytest, ) +def _account(*, account_id: str = "user-1", tenant_id: str = "tenant-1") -> Account: + account = Account(name="User", email=f"{account_id}@example.com") + account.id = account_id + tenant = Tenant(name="Tenant") + tenant.id = tenant_id + account._current_tenant = tenant + return account + + +def _dataset_row( + *, + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + built_in_field_enabled: bool = False, + data_source_type: str | None = None, + indexing_technique: str | None = "economy", +) -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Dataset", + description="", + provider="vendor", + created_by="user-1", + maintainer="user-1", + built_in_field_enabled=built_in_field_enabled, + chunk_structure=IndexStructureType.PARAGRAPH_INDEX, + data_source_type=data_source_type, + indexing_technique=indexing_technique, + ) + + +def _document_row( + *, + document_id: str = "document-1", + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + name: str = "Document", + indexing_status: str = IndexingStatus.COMPLETED, + data_source_type: str = DataSourceType.UPLOAD_FILE, + data_source_info: str = "{}", + enabled: bool = True, + archived: bool = False, + is_paused: bool = False, +) -> Document: + return Document( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=data_source_type, + data_source_info=data_source_info, + batch="batch-1", + name=name, + created_from=DocumentCreatedFrom.API, + created_by="user-1", + created_at=datetime(2026, 1, 1), + updated_at=datetime(2026, 1, 2), + indexing_status=indexing_status, + doc_form=IndexStructureType.PARAGRAPH_INDEX, + word_count=10, + enabled=enabled, + archived=archived, + is_paused=is_paused, + completed_at=datetime(2026, 1, 2) if indexing_status == IndexingStatus.COMPLETED else None, + ) + + +def _upload_file(*, file_id: str, tenant_id: str = "tenant-1", name: str = "upload.txt") -> UploadFile: + upload_file = UploadFile( + tenant_id=tenant_id, + storage_type="opendal", + key=f"key-{file_id}", + name=name, + size=1, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="user-1", + created_at=datetime(2026, 1, 1), + used=False, + ) + upload_file.id = file_id + return upload_file + + +def _process_rule( + *, dataset_id: str = "dataset-1", rule_id: str = "rule-1", mode: str = "automatic" +) -> DatasetProcessRule: + rules = ( + json.dumps(DatasetProcessRule.AUTOMATIC_RULES) + if mode == "automatic" + else Rule( + pre_processing_rules=[PreProcessingRule(id="remove_extra_spaces", enabled=True)], + segmentation=Segmentation(separator="\n", max_tokens=100), + ).model_dump_json() + ) + process_rule = DatasetProcessRule(dataset_id=dataset_id, mode=mode, rules=rules, created_by="user-1") + process_rule.id = rule_id + return process_rule + + class _RetryFlagLock: def __init__(self, store: "_RetryFlagStore", key: str): self.store = store @@ -102,183 +209,212 @@ class TestDocumentServiceDisplayStatus: assert DocumentService.build_display_status_filters("missing") == () def test_apply_display_status_filter_returns_original_query_for_unknown_status(self): - query = MagicMock() + query = select(Document) result = DocumentService.apply_display_status_filter(query, "missing") assert result is query - query.where.assert_not_called() def test_apply_display_status_filter_applies_where_for_known_status(self): - query = MagicMock() - filtered_query = MagicMock() - query.where.return_value = filtered_query + query = select(Document) result = DocumentService.apply_display_status_filter(query, "enabled") - assert result is filtered_query - query.where.assert_called_once() + assert result is not query + assert "documents.enabled" in str(result) class TestDocumentServiceRetrieval: - def test_get_document_by_id_uses_provided_session(self): - session = MagicMock() - expected_document = DatasetServiceUnitDataFactory.create_document_mock(document_id="document-1") - session.get.return_value = expected_document + def test_get_document_by_id_uses_provided_session(self, sqlite_session: Session): + document = _document_row() + sqlite_session.add(document) + sqlite_session.commit() - result = DocumentService.get_document_by_id("document-1", session=session) + assert DocumentService.get_document_by_id(document.id, session=sqlite_session) is document + assert DocumentService.get_document_by_id("missing", session=sqlite_session) is None - assert result is expected_document - session.get.assert_called_once_with(Document, "document-1") + def test_get_document_by_ids_enforces_dataset_owner_and_state(self, sqlite_session: Session): + dataset = _dataset_row() + expected = _document_row(document_id="expected") + sqlite_session.add_all( + [ + dataset, + expected, + _document_row(document_id="disabled", enabled=False), + _document_row(document_id="archived", archived=True), + _document_row(document_id="waiting", indexing_status=IndexingStatus.WAITING), + _document_row(document_id="other-dataset", dataset_id="dataset-2"), + _document_row(document_id="other-tenant", tenant_id="tenant-2"), + ] + ) + sqlite_session.commit() + + documents = DocumentService.get_document_by_ids( + DatasetRefService.create_dataset_ref(dataset), + ["expected", "disabled", "archived", "waiting", "other-dataset", "other-tenant"], + sqlite_session, + ) + + assert [document.id for document in documents] == [expected.id] class TestDocumentServiceMutations: """Unit tests for DocumentService mutation and orchestration helpers.""" - @pytest.fixture - def rename_account_context(self): - class FakeAccount: - pass - - current_user = FakeAccount() - current_user.id = "user-123" - current_user.current_tenant_id = "tenant-123" - - with ( - patch("services.dataset_service.Account", FakeAccount), - patch("services.dataset_service.current_user", current_user), - ): - yield current_user - @pytest.mark.parametrize(("archived", "expected"), [(True, True), (False, False)]) def test_check_archived_returns_boolean_status(self, archived, expected): - document = DatasetServiceUnitDataFactory.create_document_mock(archived=archived) + document = _document_row(archived=archived) assert DocumentService.check_archived(document) is expected - def test_delete_documents_limits_query_and_cleanup_to_dataset_ref(self): - session = MagicMock() - dataset = _make_dataset( - dataset_id="dataset-1", - tenant_id="tenant-1", - doc_form=IndexStructureType.PARAGRAPH_INDEX, + def test_delete_documents_limits_query_and_cleanup_to_dataset_ref(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row( + document_id="doc-1", + data_source_info=json.dumps({"upload_file_id": "file-1"}), ) - document = _make_document(document_id="doc-1", dataset_id=dataset.id, tenant_id=dataset.tenant_id) - - with ( - patch("services.dataset_service.batch_clean_document_task") as clean_task, - ): - session.scalars.return_value.all.return_value = [document] + other_dataset = _document_row(document_id="other-dataset", dataset_id="dataset-2") + other_tenant = _document_row(document_id="other-tenant", tenant_id="tenant-2") + sqlite_session.add_all([dataset, document, other_dataset, other_tenant]) + sqlite_session.commit() + with patch("services.dataset_service.batch_clean_document_task") as clean_task: dataset_ref = DatasetRefService.create_dataset_ref(dataset) DocumentService.delete_documents( dataset_ref, - ["doc-1", "other-doc"], - dataset.doc_form, - session, + [document.id, other_dataset.id, other_tenant.id], + IndexStructureType.PARAGRAPH_INDEX, + sqlite_session, ) - stmt = session.scalars.call_args.args[0] - compiled = stmt.compile() - statement = str(compiled) - assert "documents.id IN" in statement - assert "documents.tenant_id" in statement - assert "documents.dataset_id" in statement - assert ["doc-1", "other-doc"] in compiled.params.values() - assert dataset.tenant_id in compiled.params.values() - assert dataset.id in compiled.params.values() - session.delete.assert_called_once_with(document) - session.commit.assert_called_once() - clean_task.delay.assert_called_once_with(["doc-1"], dataset.id, dataset.doc_form, []) + assert sqlite_session.get(Document, document.id) is None + assert sqlite_session.get(Document, other_dataset.id) is other_dataset + assert sqlite_session.get(Document, other_tenant.id) is other_tenant + clean_task.delay.assert_called_once_with( + [document.id], dataset.id, IndexStructureType.PARAGRAPH_INDEX, ["file-1"] + ) - def test_rename_document_raises_when_dataset_is_missing(self, rename_account_context): - session = MagicMock() + def test_delete_documents_with_empty_ids_does_not_commit(self, sqlite_session: Session): + commits = 0 - with patch.object(DatasetService, "get_dataset", return_value=None): + def count_commit(_session): + nonlocal commits + commits += 1 + + event.listen(sqlite_session, "after_commit", count_commit) + DocumentService.delete_documents( + DatasetRefService.create_dataset_ref(_dataset_row()), [], IndexStructureType.PARAGRAPH_INDEX, sqlite_session + ) + event.remove(sqlite_session, "after_commit", count_commit) + assert commits == 0 + + def test_rename_document_raises_when_dataset_is_missing(self, sqlite_session: Session): + with patch("services.dataset_service.current_user", _account()): with pytest.raises(ValueError, match="Dataset not found"): - DocumentService.rename_document("dataset-1", "doc-1", "New Name", session) + DocumentService.rename_document("dataset-1", "doc-1", "New Name", sqlite_session) - def test_rename_document_raises_when_document_is_missing(self, rename_account_context): - dataset = DatasetServiceUnitDataFactory.create_dataset_mock() - session = MagicMock() - - with ( - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DocumentService, "get_document", return_value=None), - ): + def test_rename_document_raises_when_document_is_missing(self, sqlite_session: Session): + dataset = _dataset_row() + sqlite_session.add(dataset) + sqlite_session.commit() + with patch("services.dataset_service.current_user", _account()): with pytest.raises(ValueError, match="Document not found"): - DocumentService.rename_document(dataset.id, "doc-1", "New Name", session) + DocumentService.rename_document(dataset.id, "doc-1", "New Name", sqlite_session) - def test_rename_document_rejects_cross_tenant_access(self, rename_account_context): - dataset = DatasetServiceUnitDataFactory.create_dataset_mock() - document = DatasetServiceUnitDataFactory.create_document_mock(tenant_id="tenant-other") - session = MagicMock() - - with ( - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DocumentService, "get_document", return_value=document), - ): + def test_rename_document_rejects_cross_tenant_access(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(tenant_id="tenant-other") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + with patch("services.dataset_service.current_user", _account()): with pytest.raises(ValueError, match="No permission"): - DocumentService.rename_document(dataset.id, document.id, "New Name", session) + DocumentService.rename_document(dataset.id, document.id, "New Name", sqlite_session) - def test_rename_document_updates_document_metadata_and_upload_file_name(self, rename_account_context): - session = MagicMock() - dataset = DatasetServiceUnitDataFactory.create_dataset_mock( - built_in_field_enabled=True, - tenant_id="tenant-1", + def test_rename_document_updates_document_metadata_and_upload_file_name(self, sqlite_session: Session): + dataset = _dataset_row(built_in_field_enabled=True) + document = _document_row( + data_source_info=json.dumps({"upload_file_id": "file-1"}), ) - document = DatasetServiceUnitDataFactory.create_document_mock( - tenant_id="tenant-1", - doc_metadata={"title": "Old"}, - data_source_info_dict={"upload_file_id": "file-1"}, + document.doc_metadata = {BuiltInField.document_name: "Old"} + upload_file = UploadFile( + tenant_id=dataset.tenant_id, + storage_type="opendal", + key="key", + name="old.txt", + size=1, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="user-1", + created_at=datetime(2026, 1, 1), + used=False, ) - rename_account_context.current_tenant_id = "tenant-1" + upload_file.id = "file-1" + sqlite_session.add_all([dataset, document, upload_file]) + sqlite_session.commit() + commits = 0 - with ( - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DocumentService, "get_document", return_value=document), - ): - result = DocumentService.rename_document(dataset.id, document.id, "New Name", session) + def count_commit(_session): + nonlocal commits + commits += 1 + + event.listen(sqlite_session, "after_commit", count_commit) + with patch("services.dataset_service.current_user", _account()): + result = DocumentService.rename_document(dataset.id, document.id, "New Name", sqlite_session) + event.remove(sqlite_session, "after_commit", count_commit) assert result is document assert document.name == "New Name" assert document.doc_metadata[BuiltInField.document_name] == "New Name" - session.add.assert_called_once_with(document) - session.execute.assert_called() - session.flush.assert_called_once() - - def test_recover_document_raises_when_document_is_not_paused(self): - document = DatasetServiceUnitDataFactory.create_document_mock(is_paused=False) - session = MagicMock() + assert sqlite_session.get(UploadFile, upload_file.id).name == "New Name" + assert commits == 0 + def test_recover_document_raises_when_document_is_not_paused(self, unbound_session: Session): + document = _document_row(is_paused=False) with pytest.raises(DocumentIndexingError): - DocumentService.recover_document(document, session) + DocumentService.recover_document(document, unbound_session) - def test_retry_document_raises_when_retry_flag_is_already_set(self, rename_account_context): - document = DatasetServiceUnitDataFactory.create_document_mock(document_id="doc-1") - session = MagicMock() + def test_recover_document_persists_and_dispatches(self, sqlite_session: Session): + document = _document_row(is_paused=True) + sqlite_session.add(document) + sqlite_session.commit() + with ( + patch("services.dataset_service.redis_client") as redis, + patch("services.dataset_service.recover_document_indexing_task") as task, + ): + DocumentService.recover_document(document, sqlite_session) - with patch("services.dataset_service.redis_client") as mock_redis: - mock_redis.lock.return_value.acquire.return_value = False + sqlite_session.expire_all() + recovered = sqlite_session.get(Document, document.id) + assert recovered is not None + assert recovered.is_paused is False + redis.delete.assert_called_once_with(f"document_{document.id}_is_paused") + task.delay.assert_called_once_with(document.dataset_id, document.id) + def test_retry_document_raises_when_retry_flag_is_already_set(self, sqlite_session: Session): + document = _document_row(indexing_status=IndexingStatus.ERROR) + sqlite_session.add(document) + sqlite_session.commit() + retry_flags = _RetryFlagStore({f"document_{document.id}_is_retried": "other-request"}) + with ( + patch("services.dataset_service.current_user", _account()), + patch("services.dataset_service.redis_client", retry_flags), + ): with pytest.raises(ValueError, match="being retried"): - DocumentService.retry_document("dataset-1", [document], session) + DocumentService.retry_document("dataset-1", [document], sqlite_session) def test_retry_document_leaves_batch_unchanged_when_later_document_is_already_being_retried( - self, rename_account_context + self, sqlite_session: Session ): - first_document = DatasetServiceUnitDataFactory.create_document_mock( - document_id="doc-1", indexing_status="error" - ) - second_document = DatasetServiceUnitDataFactory.create_document_mock( - document_id="doc-2", indexing_status="error" - ) + first_document = _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR) + second_document = _document_row(document_id="doc-2", indexing_status=IndexingStatus.ERROR) + sqlite_session.add_all([first_document, second_document]) + sqlite_session.commit() first_retry_key = "document_doc-1_is_retried" second_retry_key = "document_doc-2_is_retried" retry_flags = _RetryFlagStore({second_retry_key: "other-request"}) - session = MagicMock() - with ( + patch("services.dataset_service.current_user", _account()), patch("services.dataset_service.redis_client", retry_flags), patch("services.dataset_service.retry_document_indexing_task") as retry_task, ): @@ -286,16 +422,16 @@ class TestDocumentServiceMutations: DocumentService.retry_document( "dataset-1", [first_document, second_document], - session, + sqlite_session, ) - assert first_document.indexing_status == "error" - assert second_document.indexing_status == "error" + assert first_document.indexing_status == IndexingStatus.ERROR + assert second_document.indexing_status == IndexingStatus.ERROR assert first_retry_key not in retry_flags.values assert retry_flags.values[second_retry_key] == "other-request" retry_task.delay.assert_not_called() - def test_retry_document_does_not_release_a_retry_flag_reacquired_by_another_request(self, rename_account_context): + def test_retry_document_does_not_release_a_retry_flag_reacquired_by_another_request(self, sqlite_session: Session): first_retry_key = "document_doc-1_is_retried" second_retry_key = "document_doc-2_is_retried" retry_flags = _RetryFlagStore( @@ -303,65 +439,93 @@ class TestDocumentServiceMutations: replacement_on_conflict=(first_retry_key, "new-owner"), ) documents = [ - DatasetServiceUnitDataFactory.create_document_mock(document_id="doc-1", indexing_status="error"), - DatasetServiceUnitDataFactory.create_document_mock(document_id="doc-2", indexing_status="error"), + _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR), + _document_row(document_id="doc-2", indexing_status=IndexingStatus.ERROR), ] + sqlite_session.add_all(documents) + sqlite_session.commit() - with patch("services.dataset_service.redis_client", retry_flags): + with ( + patch("services.dataset_service.current_user", _account()), + patch("services.dataset_service.redis_client", retry_flags), + ): with pytest.raises(ValueError, match="being retried"): - DocumentService.retry_document("dataset-1", documents, MagicMock()) + DocumentService.retry_document("dataset-1", documents, sqlite_session) assert retry_flags.values[first_retry_key] == "new-owner" assert retry_flags.values[second_retry_key] == "other-request" - def test_retry_document_releases_flags_when_status_commit_fails(self, rename_account_context): + def test_retry_document_releases_flags_when_status_commit_fails(self, sqlite_session: Session): retry_flags = _RetryFlagStore() - document = DatasetServiceUnitDataFactory.create_document_mock(document_id="doc-1", indexing_status="error") - session = MagicMock() - session.commit.side_effect = RuntimeError("database unavailable") + document = _document_row(indexing_status=IndexingStatus.ERROR) + sqlite_session.add(document) + sqlite_session.commit() + def fail_commit(_session): + raise RuntimeError("database unavailable") + + event.listen(sqlite_session, "before_commit", fail_commit) with ( + patch("services.dataset_service.current_user", _account()), patch("services.dataset_service.redis_client", retry_flags), patch("services.dataset_service.retry_document_indexing_task") as retry_task, ): with pytest.raises(RuntimeError, match="database unavailable"): - DocumentService.retry_document("dataset-1", [document], session) + DocumentService.retry_document("dataset-1", [document], sqlite_session) + event.remove(sqlite_session, "before_commit", fail_commit) assert retry_flags.values == {} - session.rollback.assert_called_once_with() retry_task.delay.assert_not_called() - def test_sync_website_document_raises_when_sync_flag_exists(self): - dataset = DatasetServiceUnitDataFactory.create_dataset_mock(dataset_id="dataset-1") - document = DatasetServiceUnitDataFactory.create_document_mock(document_id="doc-1", dataset_id=dataset.id) - session = MagicMock() + def test_retry_document_persists_status_and_dispatches(self, sqlite_session: Session): + documents = [ + _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR), + _document_row(document_id="doc-2", indexing_status=IndexingStatus.PAUSED), + ] + sqlite_session.add_all(documents) + sqlite_session.commit() + retry_flags = _RetryFlagStore() + with ( + patch("services.dataset_service.current_user", _account()), + patch("services.dataset_service.redis_client", retry_flags), + patch("services.dataset_service.retry_document_indexing_task") as task, + ): + DocumentService.retry_document("dataset-1", documents, sqlite_session) + sqlite_session.expire_all() + statuses = sqlite_session.scalars(select(Document.indexing_status).order_by(Document.id)).all() + assert statuses == [IndexingStatus.WAITING, IndexingStatus.WAITING] + task.delay.assert_called_once_with("dataset-1", ["doc-1", "doc-2"], "user-1") + + def test_sync_website_document_raises_when_sync_flag_exists(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row() with patch("services.dataset_service.redis_client") as mock_redis: mock_redis.get.return_value = "1" with pytest.raises(ValueError, match="being synced"): - DocumentService.sync_website_document(dataset, document, session) + DocumentService.sync_website_document(dataset, document, sqlite_session) - def test_sync_website_document_rejects_document_outside_dataset(self): - dataset = DatasetServiceUnitDataFactory.create_dataset_mock(dataset_id="dataset-1") - document = DatasetServiceUnitDataFactory.create_document_mock(document_id="doc-1", dataset_id="dataset-2") + def test_sync_website_document_rejects_document_outside_dataset(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(dataset_id="dataset-2") with ( pytest.raises(ValueError, match="Document not found"), patch("services.dataset_service.redis_client") as mock_redis, ): - DocumentService.sync_website_document(dataset, document, MagicMock()) + DocumentService.sync_website_document(dataset, document, sqlite_session) mock_redis.get.assert_not_called() - def test_sync_website_document_updates_status_sets_cache_and_dispatches_task(self): - session = MagicMock() - dataset = DatasetServiceUnitDataFactory.create_dataset_mock(dataset_id="dataset-1") - document = DatasetServiceUnitDataFactory.create_document_mock( - document_id="doc-1", - dataset_id=dataset.id, - data_source_info_dict={"mode": "crawl"}, + def test_sync_website_document_updates_status_sets_cache_and_dispatches_task(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row( + data_source_type=DataSourceType.WEBSITE_CRAWL, + data_source_info=json.dumps({"mode": "crawl"}), ) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() with ( patch("services.dataset_service.redis_client") as mock_redis, @@ -369,13 +533,14 @@ class TestDocumentServiceMutations: ): mock_redis.get.return_value = None - DocumentService.sync_website_document(dataset, document, session) + DocumentService.sync_website_document(dataset, document, sqlite_session) - assert document.indexing_status == "waiting" - assert '"mode": "scrape"' in document.data_source_info - session.add.assert_called_once_with(document) - session.commit.assert_called_once() - mock_redis.setex.assert_called_once_with("document_doc-1_is_sync", 600, 1) + sqlite_session.expire_all() + synced = sqlite_session.get(Document, document.id) + assert synced is not None + assert synced.indexing_status == IndexingStatus.WAITING + assert synced.data_source_info_dict["mode"] == "scrape" + mock_redis.setex.assert_called_once_with(f"document_{document.id}_is_sync", 600, 1) sync_task.delay.assert_called_once_with(dataset.id, document.id) @@ -384,17 +549,14 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: @pytest.fixture def account_context(self): - account = create_autospec(Account, instance=True) - account.id = "user-1" - account.current_tenant_id = "tenant-1" + account = _account() with patch("services.dataset_service.current_user", account): yield account def test_save_document_without_dataset_id_creates_high_quality_dataset_with_default_retrieval_model( - self, account_context + self, account_context, sqlite_session: Session ): - session = MagicMock() knowledge_config = KnowledgeConfig( indexing_technique="high_quality", data_source=DataSource( @@ -408,24 +570,21 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: summary_index_setting={"enable": True}, is_multimodal=True, ) - created_dataset = SimpleNamespace( - id="dataset-1", - tenant_id="tenant-1", - name="", - description=None, + binding = DatasetCollectionBinding( + provider_name="provider", + model_name="embedding-model", + type="dataset", + collection_name="collection", ) - first_document = SimpleNamespace(name="VeryLongDocumentNameForDataset.txt") + binding.id = "binding-1" + first_document = _document_row(name="VeryLongDocumentNameForDataset.txt") with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch( "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", - return_value=SimpleNamespace(id="binding-1"), + return_value=binding, ), - patch( - "services.dataset_service.Dataset", - side_effect=lambda **kwargs: created_dataset.__dict__.update(kwargs) or created_dataset, - ) as dataset_cls, patch.object( DocumentService, "save_document_with_dataset_id", return_value=([first_document], "batch-1") ) as save_document, @@ -434,33 +593,32 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: tenant_id="tenant-1", knowledge_config=knowledge_config, account=account_context, - session=session, + session=sqlite_session, ) - assert dataset is created_dataset assert documents == [first_document] assert batch == "batch-1" - assert created_dataset.collection_binding_id == "binding-1" - assert created_dataset.retrieval_model["search_method"] == RetrievalMethod.SEMANTIC_SEARCH - assert created_dataset.retrieval_model["top_k"] == 4 - assert created_dataset.summary_index_setting == {"enable": True} - assert created_dataset.is_multimodal is True - assert created_dataset.name == first_document.name[:18] + "..." + assert dataset.collection_binding_id == "binding-1" + assert dataset.retrieval_model["search_method"] == RetrievalMethod.SEMANTIC_SEARCH + assert dataset.retrieval_model["top_k"] == 4 + assert dataset.summary_index_setting == {"enable": True} + assert dataset.is_multimodal is True + assert dataset.name == first_document.name[:18] + "..." assert ( - created_dataset.description + dataset.description == "useful for when you want to answer queries about the VeryLongDocumentNameForDataset.txt" ) - dataset_cls.assert_called_once() + assert sqlite_session.get(Dataset, dataset.id) is dataset save_document.assert_called_once_with( - created_dataset, + dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) - assert session.flush.call_count == 2 - def test_save_document_without_dataset_id_uses_provided_retrieval_model(self, account_context): - session = MagicMock() + def test_save_document_without_dataset_id_uses_provided_retrieval_model( + self, account_context, sqlite_session: Session + ): retrieval_model = RetrievalModel( search_method=RetrievalMethod.SEMANTIC_SEARCH, reranking_enable=True, @@ -482,31 +640,30 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: ), retrieval_model=retrieval_model, ) - created_dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1", name="", description=None) + first_document = _document_row(name="Doc") with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), - patch( - "services.dataset_service.Dataset", - side_effect=lambda **kwargs: created_dataset.__dict__.update(kwargs) or created_dataset, - ), patch.object( DocumentService, "save_document_with_dataset_id", - return_value=([SimpleNamespace(name="Doc")], "batch-1"), + return_value=([first_document], "batch-1"), ), ): - DocumentService.save_document_without_dataset_id( + dataset, _, _ = DocumentService.save_document_without_dataset_id( "tenant-1", knowledge_config, account_context, - session, + sqlite_session, ) - assert created_dataset.retrieval_model == retrieval_model.model_dump() - assert created_dataset.collection_binding_id is None + assert dataset.retrieval_model == retrieval_model.model_dump() + assert dataset.collection_binding_id is None + assert sqlite_session.get(Dataset, dataset.id) is dataset - def test_save_document_without_dataset_id_rejects_sandbox_batch_upload(self, account_context): + def test_save_document_without_dataset_id_rejects_sandbox_batch_upload( + self, account_context, unbound_session: Session + ): knowledge_config = KnowledgeConfig( indexing_technique="economy", data_source=DataSource( @@ -524,9 +681,10 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: ), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): - session = MagicMock() with pytest.raises(ValueError, match="does not support batch upload"): - DocumentService.save_document_without_dataset_id("tenant-1", knowledge_config, account_context, session) + DocumentService.save_document_without_dataset_id( + "tenant-1", knowledge_config, account_context, unbound_session + ) check_quota.assert_not_called() @@ -536,15 +694,15 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: @pytest.fixture def account_context(self): - account = create_autospec(Account, instance=True) - account.id = "user-1" - account.current_tenant_id = "tenant-1" + account = _account() with patch("services.dataset_service.current_user", account): yield account - def test_update_document_with_dataset_id_raises_when_document_is_missing(self, account_context): - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") + def test_update_document_with_dataset_id_raises_when_document_is_missing( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -555,25 +713,24 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: ) ), ) - session = MagicMock() - - with ( - patch.object(DocumentService, "get_document", return_value=None), - patch.object(DatasetService, "check_dataset_model_setting") as check_model_setting, - ): + with patch.object(DatasetService, "check_dataset_model_setting") as check_model_setting: with pytest.raises(NotFound, match="Document not found"): DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) check_model_setting.assert_called_once_with(dataset) - def test_update_document_with_dataset_id_rejects_non_available_documents(self, account_context): - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") - document = SimpleNamespace(display_status="indexing") + def test_update_document_with_dataset_id_rejects_non_available_documents( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1", indexing_status=IndexingStatus.INDEXING) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -584,25 +741,46 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: ) ), ) - session = MagicMock() - - with ( - patch.object(DocumentService, "get_document", return_value=document), - patch.object(DatasetService, "check_dataset_model_setting"), - ): + with patch.object(DatasetService, "check_dataset_model_setting"): with pytest.raises(ValueError, match="Document is not available"): DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) - def test_update_document_with_dataset_id_upload_file_process_rule_and_name_override(self, account_context): - session = MagicMock() - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") - document = _make_document() - document.dataset_process_rule_id = "old-rule" + def test_update_document_with_dataset_id_upload_file_process_rule_and_name_override( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + upload_file = UploadFile( + tenant_id=dataset.tenant_id, + storage_type="opendal", + key="key", + name="upload.txt", + size=1, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by=account_context.id, + created_at=datetime(2026, 1, 1), + used=False, + ) + upload_file.id = "file-1" + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + position=1, + content="content", + word_count=1, + tokens=1, + created_by=account_context.id, + ) + sqlite_session.add_all([dataset, document, upload_file, segment]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -622,26 +800,23 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: name="Renamed document", doc_form=IndexStructureType.QA_INDEX, ) - created_process_rule = SimpleNamespace(id="rule-2") + updated_at = datetime(2026, 2, 1) with ( - patch.object(DocumentService, "get_document", return_value=document), patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.DatasetProcessRule", return_value=created_process_rule), - patch("services.dataset_service.naive_utc_now", return_value="now"), + patch("services.dataset_service.naive_utc_now", return_value=updated_at), patch("services.dataset_service.document_indexing_update_task") as update_task, ): - session.scalar.return_value = SimpleNamespace(id="file-1", name="upload.txt") - result = DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) assert result is document - assert document.dataset_process_rule_id == "rule-2" + assert document.dataset_process_rule_id is not None + assert sqlite_session.get(DatasetProcessRule, document.dataset_process_rule_id) is not None assert document.data_source_type == "upload_file" assert document.data_source_info == '{"upload_file_id": "file-1"}' assert document.name == "Renamed document" @@ -651,17 +826,23 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: assert document.parsing_completed_at is None assert document.cleaning_completed_at is None assert document.splitting_completed_at is None - assert document.updated_at == "now" + assert document.updated_at == updated_at assert document.created_from == "web" assert document.doc_form == IndexStructureType.QA_INDEX - assert session.commit.call_count == 3 - session.execute.assert_called() + sqlite_session.expire_all() + persisted = sqlite_session.get(Document, document.id) + assert persisted is not None + assert persisted.name == "Renamed document" + assert sqlite_session.get(DocumentSegment, segment.id).status == "re_segment" update_task.delay.assert_called_once_with(document.dataset_id, document.id) - def test_update_document_with_dataset_id_notion_import_requires_binding(self, account_context): - session = MagicMock() - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") - document = SimpleNamespace(display_status="available", id="doc-1", dataset_id="dataset-1") + def test_update_document_with_dataset_id_notion_import_requires_binding( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -679,24 +860,32 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: ), ) - with ( - patch.object(DocumentService, "get_document", return_value=document), - patch.object(DatasetService, "check_dataset_model_setting"), - ): - session.scalar.return_value = None - + with patch.object(DatasetService, "check_dataset_model_setting"): with pytest.raises(ValueError, match="Data source binding not found"): DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) - def test_update_document_with_dataset_id_website_crawl_updates_segments_and_dispatches_task(self, account_context): - session = MagicMock() - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") - document = _make_document() + def test_update_document_with_dataset_id_website_crawl_updates_segments_and_dispatches_task( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + position=1, + content="content", + word_count=1, + tokens=1, + created_by=account_context.id, + ) + sqlite_session.add_all([dataset, document, segment]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -715,16 +904,15 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: ) with ( - patch.object(DocumentService, "get_document", return_value=document), patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.naive_utc_now", return_value="now"), + patch("services.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), patch("services.dataset_service.document_indexing_update_task") as update_task, ): result = DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) assert result is document @@ -735,7 +923,8 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: ) assert document.name == "" assert document.doc_form == IndexStructureType.PARENT_CHILD_INDEX - session.execute.assert_called() + sqlite_session.expire_all() + assert sqlite_session.get(DocumentSegment, segment.id).status == "re_segment" update_task.delay.assert_called_once_with("dataset-1", "doc-1") @@ -882,9 +1071,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: @pytest.fixture def account_context(self): - account = create_autospec(Account, instance=True) - account.id = "user-1" - account.current_tenant_id = "tenant-1" + account = _account() with ( patch("services.dataset_service.current_user", account), @@ -892,22 +1079,25 @@ class TestDocumentServiceSaveDocumentWithDatasetId: ): yield account - def test_save_document_with_dataset_id_requires_file_info_for_upload_source(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_requires_file_info_for_upload_source( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(file_ids=None) with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=True)): - session = MagicMock() with pytest.raises(ValueError, match="File source info is required"): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=unbound_session, ) - def test_save_document_with_dataset_id_blocks_batch_upload_for_sandbox_plan(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_blocks_batch_upload_for_sandbox_plan( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) with ( @@ -917,19 +1107,18 @@ class TestDocumentServiceSaveDocumentWithDatasetId: ), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): - session = MagicMock() with pytest.raises(ValueError, match="does not support batch upload"): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=unbound_session, ) check_quota.assert_not_called() - def test_save_document_with_dataset_id_enforces_batch_upload_limit(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_enforces_batch_upload_limit(self, account_context, unbound_session: Session): + dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) with ( @@ -937,21 +1126,23 @@ class TestDocumentServiceSaveDocumentWithDatasetId: patch("services.dataset_service.dify_config.BATCH_UPLOAD_LIMIT", 1), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): - session = MagicMock() with pytest.raises(ValueError, match="batch upload limit of 1"): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=unbound_session, ) check_quota.assert_not_called() - def test_save_document_with_dataset_id_updates_existing_document_and_data_source_type(self, account_context): - dataset = _make_dataset(data_source_type=None) + def test_save_document_with_dataset_id_updates_existing_document_and_data_source_type( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=None) knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) - updated_document = _make_document(document_id="doc-1", batch="batch-existing") + updated_document = _document_row(document_id="doc-1") + updated_document.batch = "batch-existing" with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), @@ -959,55 +1150,57 @@ class TestDocumentServiceSaveDocumentWithDatasetId: DocumentService, "update_document_with_dataset_id", return_value=updated_document ) as update_document, ): - session = MagicMock() documents, batch = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) assert dataset.data_source_type == "upload_file" assert documents == [updated_document] assert batch == "batch-existing" - update_document.assert_called_once_with(dataset, knowledge_config, account_context, session=session) + update_document.assert_called_once_with(dataset, knowledge_config, account_context, session=sqlite_session) - def test_save_document_with_dataset_id_requires_data_source_for_new_documents(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_requires_data_source_for_new_documents( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(data_source=None) with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)): - session = MagicMock() with pytest.raises(ValueError, match="Data source is required when creating new documents"): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=unbound_session, ) - def test_save_document_with_dataset_id_requires_existing_process_rule_for_custom_mode(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_requires_existing_process_rule_for_custom_mode( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + sqlite_session.add(dataset) + sqlite_session.commit() knowledge_config = _make_upload_knowledge_config( file_ids=["file-1"], process_rule=ProcessRule(mode="custom"), ) with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)): - session = MagicMock() - session.scalar.return_value = None with pytest.raises(ValueError, match="No process rule found"): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) - session.scalar.assert_called_once() - - def test_save_document_with_dataset_id_rejects_invalid_indexing_technique(self, account_context): - dataset = _make_dataset(indexing_technique=None) + def test_save_document_with_dataset_id_rejects_invalid_indexing_technique( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row(indexing_technique=None) knowledge_config = SimpleNamespace( doc_form=IndexStructureType.PARAGRAPH_INDEX, original_document_id=None, @@ -1016,17 +1209,18 @@ class TestDocumentServiceSaveDocumentWithDatasetId: ) with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)): - session = MagicMock() with pytest.raises(ValueError, match="Indexing technique is invalid"): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=unbound_session, ) - def test_save_document_with_dataset_id_returns_empty_for_invalid_process_rule_mode(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_returns_empty_for_invalid_process_rule_mode( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(file_ids=["file-1"]) knowledge_config.process_rule = SimpleNamespace(mode="unsupported-mode", rules=None) @@ -1035,77 +1229,61 @@ class TestDocumentServiceSaveDocumentWithDatasetId: dataset, knowledge_config, account_context, - session=MagicMock(), + session=unbound_session, ) assert documents == [] assert batch == "" - def test_save_document_with_dataset_id_upload_file_creates_and_reindexes_documents(self, account_context): - session = MagicMock() - dataset = _make_dataset() - dataset_process_rule = SimpleNamespace(id="rule-1") + def test_save_document_with_dataset_id_upload_file_creates_and_reindexes_documents( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) + dataset_process_rule = _process_rule() knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) - duplicate_document = _make_document(document_id="doc-duplicate", name="existing.txt") - created_document = _make_document(document_id="doc-created", name="new.txt") - upload_file_a = SimpleNamespace(id="file-1", name="existing.txt") - upload_file_b = SimpleNamespace(id="file-2", name="new.txt") + duplicate_document = _document_row(document_id="doc-duplicate", name="existing.txt") + upload_file_a = _upload_file(file_id="file-1", name="existing.txt") + upload_file_b = _upload_file(file_id="file-2", name="new.txt") + sqlite_session.add_all([dataset, dataset_process_rule, duplicate_document, upload_file_a, upload_file_b]) + sqlite_session.commit() with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch.object(DocumentService, "get_documents_position", return_value=4), - patch.object(DocumentService, "build_document", return_value=created_document) as build_document, patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, patch("services.dataset_service.DuplicateDocumentIndexingTaskProxy") as duplicate_proxy_cls, - patch("services.dataset_service.naive_utc_now", return_value="now"), + patch("services.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), patch("services.dataset_service.time.strftime", return_value="20260101010101"), patch("services.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() - session.scalars.return_value.all.side_effect = [ - [upload_file_a, upload_file_b], - [duplicate_document], - ] - documents, batch = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, dataset_process_rule=dataset_process_rule, - session=session, + session=sqlite_session, ) - assert documents == [duplicate_document, created_document] + assert [document.name for document in documents] == ["existing.txt", "new.txt"] assert batch == "20260101010101100023" assert duplicate_document.dataset_process_rule_id == "rule-1" - assert duplicate_document.updated_at == "now" + assert duplicate_document.updated_at == datetime(2026, 2, 1) assert duplicate_document.batch == batch - assert duplicate_document.indexing_status == "waiting" - build_document.assert_called_once_with( - dataset, - "rule-1", - "upload_file", - IndexStructureType.PARAGRAPH_INDEX, - "English", - {"upload_file_id": "file-2"}, - "web", - 4, - account_context, - "new.txt", - batch, - ) - document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-created"]) + assert duplicate_document.indexing_status == IndexingStatus.WAITING + created_document = next(document for document in documents if document.name == "new.txt") + sqlite_session.expire_all() + assert sqlite_session.get(Document, created_document.id) is not None + document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, [created_document.id]) document_proxy_cls.return_value.delay.assert_called_once() duplicate_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-duplicate"]) duplicate_proxy_cls.return_value.delay.assert_called_once() def test_save_document_with_dataset_id_notion_import_truncates_names_and_cleans_removed_pages( - self, account_context + self, account_context, sqlite_session: Session ): - session = MagicMock() - dataset = _make_dataset() - dataset_process_rule = SimpleNamespace(id="rule-1") + dataset = _dataset_row(data_source_type=DataSourceType.NOTION_IMPORT) + dataset_process_rule = _process_rule() notion_page_name = "a" * 300 knowledge_config = KnowledgeConfig( indexing_technique="economy", @@ -1132,42 +1310,50 @@ class TestDocumentServiceSaveDocumentWithDatasetId: doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_language="English", ) - existing_keep = _make_document(document_id="doc-keep") + existing_keep = _document_row( + document_id="doc-keep", + data_source_type=DataSourceType.NOTION_IMPORT, + ) existing_keep.data_source_info = json.dumps({"notion_page_id": "page-keep"}) - existing_remove = _make_document(document_id="doc-remove") + existing_remove = _document_row( + document_id="doc-remove", + data_source_type=DataSourceType.NOTION_IMPORT, + ) existing_remove.data_source_info = json.dumps({"notion_page_id": "page-remove"}) - created_document = _make_document(document_id="doc-new") + sqlite_session.add_all([dataset, dataset_process_rule, existing_keep, existing_remove]) + sqlite_session.commit() with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch.object(DocumentService, "get_documents_position", return_value=1), - patch.object(DocumentService, "build_document", return_value=created_document) as build_document, patch("services.dataset_service.clean_notion_document_task") as clean_task, patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, patch("services.dataset_service.uuid.uuid4", return_value="doc-new"), ): mock_redis.lock.return_value = _make_lock_context() - session.scalars.return_value.all.return_value = [existing_keep, existing_remove] - documents, _ = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, dataset_process_rule=dataset_process_rule, - session=session, + session=sqlite_session, ) + created_document = next(document for document in documents if document.id == "doc-new") assert created_document in documents - assert len(build_document.call_args.args[9]) == 255 + assert len(created_document.name) == 255 + assert sqlite_session.get(Document, created_document.id) is created_document clean_task.delay.assert_called_once_with(["doc-remove"], dataset.id) document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-new"]) document_proxy_cls.return_value.delay.assert_called_once() - def test_save_document_with_dataset_id_website_crawl_truncates_long_urls(self, account_context): - session = MagicMock() - dataset = _make_dataset() - dataset_process_rule = SimpleNamespace(id="rule-1") + def test_save_document_with_dataset_id_website_crawl_truncates_long_urls( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.WEBSITE_CRAWL) + dataset_process_rule = _process_rule() + sqlite_session.add_all([dataset, dataset_process_rule]) + sqlite_session.commit() long_url = "https://example.com/" + ("a" * 260) short_url = "https://example.com/short" knowledge_config = KnowledgeConfig( @@ -1186,18 +1372,9 @@ class TestDocumentServiceSaveDocumentWithDatasetId: doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_language="English", ) - first_document = _make_document(document_id="doc-1") - second_document = _make_document(document_id="doc-2") - with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch.object(DocumentService, "get_documents_position", return_value=2), - patch.object( - DocumentService, - "build_document", - side_effect=[first_document, second_document], - ) as build_document, patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, ): mock_redis.lock.return_value = _make_lock_context() @@ -1207,13 +1384,14 @@ class TestDocumentServiceSaveDocumentWithDatasetId: knowledge_config, account_context, dataset_process_rule=dataset_process_rule, - session=session, + session=sqlite_session, ) - assert documents == [first_document, second_document] - assert build_document.call_args_list[0].args[9] == long_url[:200] + "..." - assert build_document.call_args_list[1].args[9] == short_url - document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-1", "doc-2"]) + assert [document.name for document in documents] == [long_url[:200] + "...", short_url] + assert sqlite_session.scalars(select(Document).order_by(Document.position)).all() == documents + document_proxy_cls.assert_called_once_with( + dataset.tenant_id, dataset.id, [document.id for document in documents] + ) document_proxy_cls.return_value.delay.assert_called_once() @@ -1221,16 +1399,16 @@ class TestDocumentServiceBatchUpdateStatus: """Unit tests for batch_update_document_status orchestration and helper branches.""" def test_prepare_disable_update_requires_completed_document(self): - document = _make_document(indexing_status="waiting") + document = _document_row(indexing_status=IndexingStatus.WAITING) document.completed_at = None with pytest.raises(DocumentIndexingError, match="is not completed"): - DocumentService._prepare_disable_update(document, user=SimpleNamespace(id="user-1"), now="now") + DocumentService._prepare_disable_update(document, user=_account(), now=datetime(2026, 2, 1)) def test_prepare_archive_update_sets_async_task_for_enabled_document(self): - document = _make_document(enabled=True, archived=False) + document = _document_row(enabled=True, archived=False) - result = DocumentService._prepare_archive_update(document, user=SimpleNamespace(id="user-1"), now="now") + result = DocumentService._prepare_archive_update(document, user=_account(), now=datetime(2026, 2, 1)) assert result is not None assert result["updates"]["archived"] is True @@ -1238,59 +1416,63 @@ class TestDocumentServiceBatchUpdateStatus: assert result["async_task"]["args"] == [document.id] def test_prepare_unarchive_update_sets_async_task_for_enabled_document(self): - document = _make_document(enabled=True, archived=True) + document = _document_row(enabled=True, archived=True) - result = DocumentService._prepare_unarchive_update(document, now="now") + result = DocumentService._prepare_unarchive_update(document, now=datetime(2026, 2, 1)) assert result is not None assert result["updates"]["archived"] is False assert result["set_cache"] is True assert result["async_task"]["args"] == [document.id] - def test_batch_update_document_status_rejects_indexing_documents(self): - session = MagicMock() - dataset = _make_dataset() - document = _make_document(name="Busy document") + def test_batch_update_document_status_rejects_indexing_documents(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(name="Busy document") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() - with ( - patch.object(DocumentService, "get_document", return_value=document), - patch("services.dataset_service.redis_client") as mock_redis, - ): + with patch("services.dataset_service.redis_client") as mock_redis: mock_redis.get.return_value = "1" with pytest.raises(DocumentIndexingError, match="Busy document is being indexed"): DocumentService.batch_update_document_status( - dataset, [document.id], "archive", SimpleNamespace(id="user-1"), session + dataset, [document.id], "archive", _account(), sqlite_session ) - session.flush.assert_not_called() + sqlite_session.refresh(document) + assert document.archived is False - def test_batch_update_document_status_rolls_back_when_commit_fails(self): - session = MagicMock() - dataset = _make_dataset() - document = _make_document(enabled=False) + def test_batch_update_document_status_rolls_back_when_commit_fails(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(enabled=False) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + def fail_commit(_session): + raise RuntimeError("commit failed") + + event.listen(sqlite_session, "before_commit", fail_commit) with ( - patch.object(DocumentService, "get_document", return_value=document), patch("services.dataset_service.redis_client") as mock_redis, ): mock_redis.get.return_value = None - session.commit.side_effect = RuntimeError("commit failed") with pytest.raises(RuntimeError, match="commit failed"): DocumentService.batch_update_document_status( - dataset, [document.id], "enable", SimpleNamespace(id="user-1"), session + dataset, [document.id], "enable", _account(), sqlite_session ) + event.remove(sqlite_session, "before_commit", fail_commit) - session.rollback.assert_called_once() + sqlite_session.refresh(document) + assert document.enabled is False - def test_batch_update_document_status_raises_async_task_error_after_commit(self): - session = MagicMock() - dataset = _make_dataset() - document = _make_document(enabled=False) + def test_batch_update_document_status_raises_async_task_error_after_commit(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(enabled=False) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() with ( - patch.object(DocumentService, "get_document", return_value=document), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.add_document_to_index_task") as add_task, ): @@ -1299,10 +1481,11 @@ class TestDocumentServiceBatchUpdateStatus: with pytest.raises(RuntimeError, match="task failed"): DocumentService.batch_update_document_status( - dataset, [document.id], "enable", SimpleNamespace(id="user-1"), session + dataset, [document.id], "enable", _account(), sqlite_session ) - session.commit.assert_called_once() + sqlite_session.refresh(document) + assert document.enabled is True mock_redis.setex.assert_called_once_with(f"document_{document.id}_indexing", 600, 1) @@ -1311,25 +1494,36 @@ class TestDocumentServiceTenantAndUpdateEdges: @pytest.fixture def account_context(self): - account = create_autospec(Account, instance=True) - account.id = "user-1" - account.current_tenant_id = "tenant-1" + account = _account() with patch("services.dataset_service.current_user", account): yield account - def test_get_tenant_documents_count_returns_query_count(self, account_context): - session = MagicMock() - session.scalar.return_value = 12 + def test_get_tenant_documents_count_scopes_state_and_tenant(self, account_context, sqlite_session: Session): + sqlite_session.add_all( + [ + _document_row(document_id="one"), + _document_row(document_id="two"), + _document_row(document_id="disabled", enabled=False), + _document_row(document_id="archived", archived=True), + _document_row(document_id="unfinished", indexing_status=IndexingStatus.WAITING), + _document_row(document_id="foreign", tenant_id="tenant-2"), + ] + ) + sqlite_session.commit() - result = DocumentService.get_tenant_documents_count(session) + result = DocumentService.get_tenant_documents_count(sqlite_session) - assert result == 12 + assert result == 2 - def test_update_document_with_dataset_id_uses_automatic_process_rule_payload(self, account_context): - session = MagicMock() - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") - document = _make_document() + def test_update_document_with_dataset_id_uses_automatic_process_rule_payload( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + upload_file = _upload_file(file_id="file-1") + sqlite_session.add_all([dataset, document, upload_file]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -1348,62 +1542,56 @@ class TestDocumentServiceTenantAndUpdateEdges: ), doc_form=IndexStructureType.PARAGRAPH_INDEX, ) - created_process_rule = SimpleNamespace(id="rule-2") + updated_at = datetime(2026, 2, 1) with ( - patch.object(DocumentService, "get_document", return_value=document), - patch("services.dataset_service.DatasetProcessRule") as process_rule_cls, patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.naive_utc_now", return_value="now"), + patch("services.dataset_service.naive_utc_now", return_value=updated_at), patch("services.dataset_service.document_indexing_update_task") as update_task, ): - process_rule_cls.AUTOMATIC_RULES = DatasetProcessRule.AUTOMATIC_RULES - process_rule_cls.return_value = created_process_rule - session.scalar.return_value = SimpleNamespace(id="file-1", name="upload.txt") - result = DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) assert result is document - assert document.dataset_process_rule_id == "rule-2" + assert document.dataset_process_rule_id is not None assert document.name == "upload.txt" - assert process_rule_cls.call_args.kwargs == { - "dataset_id": "dataset-1", - "mode": "automatic", - "rules": json.dumps(DatasetProcessRule.AUTOMATIC_RULES), - "created_by": "user-1", - } - assert session.commit.call_count == 3 + process_rule = sqlite_session.get(DatasetProcessRule, document.dataset_process_rule_id) + assert process_rule is not None + assert process_rule.mode == "automatic" + assert process_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) update_task.delay.assert_called_once_with("dataset-1", "doc-1") - def test_update_document_with_dataset_id_requires_upload_file_info(self, account_context): - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") + def test_update_document_with_dataset_id_requires_upload_file_info(self, account_context, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", data_source=DataSource(info_list=InfoList(data_source_type="upload_file")), ) - with ( - patch.object(DocumentService, "get_document", return_value=_make_document()), - patch.object(DatasetService, "check_dataset_model_setting"), - ): - session = MagicMock() + with patch.object(DatasetService, "check_dataset_model_setting"): with pytest.raises(ValueError, match="No file info list found"): DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) - def test_update_document_with_dataset_id_raises_when_upload_file_is_missing(self, account_context): - session = MagicMock() - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") + def test_update_document_with_dataset_id_raises_when_upload_file_is_missing( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -1415,45 +1603,49 @@ class TestDocumentServiceTenantAndUpdateEdges: ), ) - with ( - patch.object(DocumentService, "get_document", return_value=_make_document()), - patch.object(DatasetService, "check_dataset_model_setting"), - ): - session.scalar.return_value = None - + with patch.object(DatasetService, "check_dataset_model_setting"): with pytest.raises(FileNotExistsError): DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) - def test_update_document_with_dataset_id_requires_notion_info_list(self, account_context): - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") + def test_update_document_with_dataset_id_requires_notion_info_list(self, account_context, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", data_source=DataSource(info_list=InfoList(data_source_type="notion_import")), ) - with ( - patch.object(DocumentService, "get_document", return_value=_make_document()), - patch.object(DatasetService, "check_dataset_model_setting"), - ): - session = MagicMock() + with patch.object(DatasetService, "check_dataset_model_setting"): with pytest.raises(ValueError, match="No notion info list found"): DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) - def test_update_document_with_dataset_id_notion_import_updates_page_info(self, account_context): - session = MagicMock() - dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1") - document = _make_document() + def test_update_document_with_dataset_id_notion_import_updates_page_info( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + binding = DataSourceOauthBinding( + tenant_id=dataset.tenant_id, + access_token="token", + provider="notion", + source_info={"workspace_id": '"workspace-1"'}, + disabled=False, + ) + sqlite_session.add_all([dataset, document, binding]) + sqlite_session.commit() document_data = KnowledgeConfig( original_document_id="doc-1", indexing_technique="economy", @@ -1476,18 +1668,15 @@ class TestDocumentServiceTenantAndUpdateEdges: ) with ( - patch.object(DocumentService, "get_document", return_value=document), patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.naive_utc_now", return_value="now"), + patch("services.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), patch("services.dataset_service.document_indexing_update_task") as update_task, ): - session.scalar.return_value = SimpleNamespace(id="binding-1") - result = DocumentService.update_document_with_dataset_id( dataset, document_data, account_context, - session=session, + session=sqlite_session, ) assert result is document @@ -1502,6 +1691,7 @@ class TestDocumentServiceTenantAndUpdateEdges: "type": "database", } ) + sqlite_session.refresh(document) update_task.delay.assert_called_once_with("dataset-1", "doc-1") @@ -1510,15 +1700,14 @@ class TestDocumentServiceSaveWithoutDatasetBilling: @pytest.fixture def account_context(self): - account = create_autospec(Account, instance=True) - account.id = "user-1" - account.current_tenant_id = "tenant-1" + account = _account() with patch("services.dataset_service.current_user", account): yield account - def test_save_document_without_dataset_id_counts_notion_pages_for_quota(self, account_context): - session = MagicMock() + def test_save_document_without_dataset_id_counts_notion_pages_for_quota( + self, account_context, sqlite_session: Session + ): knowledge_config = KnowledgeConfig( indexing_technique="economy", data_source=DataSource( @@ -1542,33 +1731,32 @@ class TestDocumentServiceSaveWithoutDatasetBilling: ) ), ) - created_dataset = SimpleNamespace(id="dataset-1", tenant_id="tenant-1", name="", description=None) features = _make_features(enabled=True) + document = _document_row(name="Doc") with ( patch("services.dataset_service.FeatureService.get_features", return_value=features), patch("services.dataset_service.dify_config.BATCH_UPLOAD_LIMIT", "10"), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, - patch( - "services.dataset_service.Dataset", - side_effect=lambda **kwargs: created_dataset.__dict__.update(kwargs) or created_dataset, - ), patch.object( DocumentService, "save_document_with_dataset_id", - return_value=([SimpleNamespace(name="Doc")], "batch-1"), + return_value=([document], "batch-1"), ), ): - DocumentService.save_document_without_dataset_id( + dataset, _, _ = DocumentService.save_document_without_dataset_id( "tenant-1", knowledge_config, account_context, - session, + sqlite_session, ) check_quota.assert_called_once_with(3, features) + assert sqlite_session.get(Dataset, dataset.id) is dataset - def test_save_document_without_dataset_id_enforces_batch_limit_for_website_urls(self, account_context): + def test_save_document_without_dataset_id_enforces_batch_limit_for_website_urls( + self, account_context, unbound_session: Session + ): knowledge_config = KnowledgeConfig( indexing_technique="economy", data_source=DataSource( @@ -1589,9 +1777,10 @@ class TestDocumentServiceSaveWithoutDatasetBilling: patch("services.dataset_service.dify_config.BATCH_UPLOAD_LIMIT", "1"), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): - session = MagicMock() with pytest.raises(ValueError, match="batch upload limit of 1"): - DocumentService.save_document_without_dataset_id("tenant-1", knowledge_config, account_context, session) + DocumentService.save_document_without_dataset_id( + "tenant-1", knowledge_config, account_context, unbound_session + ) check_quota.assert_not_called() @@ -1727,9 +1916,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: @pytest.fixture def account_context(self): - account = create_autospec(Account, instance=True) - account.id = "user-1" - account.current_tenant_id = "tenant-1" + account = _account() with ( patch("services.dataset_service.current_user", account), @@ -1738,21 +1925,29 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: yield account def test_save_document_with_dataset_id_initializes_high_quality_dataset_from_default_embedding_model( - self, account_context + self, account_context, sqlite_session: Session ): - dataset = _make_dataset(data_source_type=None, indexing_technique=None) + dataset = _dataset_row(data_source_type=None, indexing_technique=None) knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) knowledge_config.indexing_technique = "high_quality" knowledge_config.embedding_model = None knowledge_config.embedding_model_provider = None - updated_document = _make_document(batch="batch-existing") + updated_document = _document_row(document_id="doc-1") + updated_document.batch = "batch-existing" + binding = DatasetCollectionBinding( + provider_name="default-provider", + model_name="default-embedding", + type="dataset", + collection_name="collection", + ) + binding.id = "binding-1" with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.ModelManager") as model_manager_cls, patch( "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", - return_value=SimpleNamespace(id="binding-1"), + return_value=binding, ) as get_binding, patch.object(DocumentService, "update_document_with_dataset_id", return_value=updated_document), ): @@ -1761,12 +1956,11 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: provider="default-provider", ) - session = MagicMock() documents, batch = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) assert documents == [updated_document] @@ -1783,10 +1977,12 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: "top_k": 4, "score_threshold_enabled": False, } - get_binding.assert_called_once_with("default-provider", "default-embedding", session) + get_binding.assert_called_once_with("default-provider", "default-embedding", sqlite_session) - def test_save_document_with_dataset_id_uses_explicit_embedding_and_retrieval_model(self, account_context): - dataset = _make_dataset(indexing_technique=None) + def test_save_document_with_dataset_id_uses_explicit_embedding_and_retrieval_model( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(indexing_technique=None) knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) knowledge_config.indexing_technique = "high_quality" knowledge_config.embedding_model = "explicit-model" @@ -1802,28 +1998,38 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: score_threshold_enabled=True, score_threshold=0.3, ) + binding = DatasetCollectionBinding( + provider_name="explicit-provider", + model_name="explicit-model", + type="dataset", + collection_name="collection", + ) + binding.id = "binding-2" + updated_document = _document_row(document_id="doc-1") with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.ModelManager") as model_manager_cls, patch( "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", - return_value=SimpleNamespace(id="binding-2"), + return_value=binding, ) as get_binding, - patch.object(DocumentService, "update_document_with_dataset_id", return_value=_make_document()), + patch.object(DocumentService, "update_document_with_dataset_id", return_value=updated_document), ): - session = MagicMock() - DocumentService.save_document_with_dataset_id(dataset, knowledge_config, account_context, session=session) + DocumentService.save_document_with_dataset_id( + dataset, knowledge_config, account_context, session=sqlite_session + ) model_manager_cls.for_tenant.return_value.get_default_model_instance.assert_not_called() - get_binding.assert_called_once_with("explicit-provider", "explicit-model", session) + get_binding.assert_called_once_with("explicit-provider", "explicit-model", sqlite_session) assert dataset.embedding_model == "explicit-model" assert dataset.embedding_model_provider == "explicit-provider" assert dataset.retrieval_model == knowledge_config.retrieval_model.model_dump() - def test_save_document_with_dataset_id_creates_custom_process_rule_for_new_upload_document(self, account_context): - session = MagicMock() - dataset = _make_dataset() + def test_save_document_with_dataset_id_creates_custom_process_rule_for_new_upload_document( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) knowledge_config = _make_upload_knowledge_config( file_ids=["file-1"], process_rule=ProcessRule( @@ -1834,148 +2040,126 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: ), ), ) - created_process_rule = SimpleNamespace(id="rule-custom") - created_document = _make_document(document_id="doc-created", name="file.txt") + upload_file = _upload_file(file_id="file-1", name="file.txt") + sqlite_session.add_all([dataset, upload_file]) + sqlite_session.commit() with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DatasetProcessRule") as process_rule_cls, - patch.object(DocumentService, "get_documents_position", return_value=3), - patch.object(DocumentService, "build_document", return_value=created_document), patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, patch("services.dataset_service.time.strftime", return_value="20260101010101"), patch("services.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() - process_rule_cls.return_value = created_process_rule - session.scalars.return_value.all.side_effect = [[SimpleNamespace(id="file-1", name="file.txt")], []] - documents, batch = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) - assert documents == [created_document] + assert len(documents) == 1 + created_document = documents[0] + assert created_document.name == "file.txt" assert batch == "20260101010101100023" - assert process_rule_cls.call_args.kwargs == { - "dataset_id": "dataset-1", - "mode": "custom", - "rules": knowledge_config.process_rule.rules.model_dump_json(), - "created_by": "user-1", - } - document_proxy_cls.assert_called_once_with("tenant-1", "dataset-1", ["doc-created"]) + created_rule = sqlite_session.get(DatasetProcessRule, created_document.dataset_process_rule_id) + assert created_rule is not None + assert created_rule.mode == "custom" + assert created_rule.rules == knowledge_config.process_rule.rules.model_dump_json() + document_proxy_cls.assert_called_once_with("tenant-1", "dataset-1", [created_document.id]) document_proxy_cls.return_value.delay.assert_called_once() def test_save_document_with_dataset_id_creates_automatic_process_rule_for_new_upload_document( - self, account_context + self, account_context, sqlite_session: Session ): - session = MagicMock() - dataset = _make_dataset() + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) knowledge_config = _make_upload_knowledge_config( file_ids=["file-1"], process_rule=ProcessRule(mode="automatic"), ) - created_process_rule = SimpleNamespace(id="rule-auto") - created_document = _make_document(document_id="doc-created", name="file.txt") + sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) + sqlite_session.commit() with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DatasetProcessRule") as process_rule_cls, - patch.object(DocumentService, "get_documents_position", return_value=1), - patch.object(DocumentService, "build_document", return_value=created_document), patch("services.dataset_service.DocumentIndexingTaskProxy"), patch("services.dataset_service.time.strftime", return_value="20260101010101"), patch("services.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() - process_rule_cls.AUTOMATIC_RULES = DatasetProcessRule.AUTOMATIC_RULES - process_rule_cls.return_value = created_process_rule - session.scalars.return_value.all.side_effect = [[SimpleNamespace(id="file-1", name="file.txt")], []] - - DocumentService.save_document_with_dataset_id( + documents, _ = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) - assert process_rule_cls.call_args.kwargs == { - "dataset_id": "dataset-1", - "mode": "automatic", - "rules": json.dumps(DatasetProcessRule.AUTOMATIC_RULES), - "created_by": "user-1", - } - assert session.flush.call_count >= 2 + created_rule = sqlite_session.get(DatasetProcessRule, documents[0].dataset_process_rule_id) + assert created_rule is not None + assert created_rule.mode == "automatic" + assert created_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) + assert sqlite_session.get(Document, documents[0].id) is documents[0] def test_save_document_with_dataset_id_creates_fallback_automatic_process_rule_when_latest_is_missing( - self, account_context + self, account_context, sqlite_session: Session ): - session = MagicMock() - dataset = _make_dataset() + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) knowledge_config = _make_upload_knowledge_config(file_ids=["file-1"], process_rule=None) - created_process_rule = SimpleNamespace(id="rule-fallback") - created_document = _make_document(document_id="doc-created", name="file.txt") + sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) + sqlite_session.commit() with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DatasetProcessRule") as process_rule_cls, - patch.object(DocumentService, "get_documents_position", return_value=1), - patch.object(DocumentService, "build_document", return_value=created_document), patch("services.dataset_service.DocumentIndexingTaskProxy"), patch("services.dataset_service.time.strftime", return_value="20260101010101"), patch("services.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() - process_rule_cls.AUTOMATIC_RULES = DatasetProcessRule.AUTOMATIC_RULES - process_rule_cls.return_value = created_process_rule - session.scalar.return_value = None - session.scalars.return_value.all.side_effect = [[SimpleNamespace(id="file-1", name="file.txt")], []] - - DocumentService.save_document_with_dataset_id( + documents, _ = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) - session.scalar.assert_called_once() - assert process_rule_cls.call_args.kwargs == { - "dataset_id": "dataset-1", - "mode": "automatic", - "rules": json.dumps(DatasetProcessRule.AUTOMATIC_RULES), - "created_by": "user-1", - } + created_rule = sqlite_session.get(DatasetProcessRule, documents[0].dataset_process_rule_id) + assert created_rule is not None + assert created_rule.mode == "automatic" + assert created_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) - def test_save_document_with_dataset_id_raises_when_upload_file_lookup_is_incomplete(self, account_context): - session = MagicMock() - dataset = _make_dataset() + def test_save_document_with_dataset_id_raises_when_upload_file_lookup_is_incomplete( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) + sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) + sqlite_session.commit() with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch.object(DocumentService, "get_documents_position", return_value=1), patch("services.dataset_service.time.strftime", return_value="20260101010101"), patch("services.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() - session.scalars.return_value.all.return_value = [SimpleNamespace(id="file-1", name="file.txt")] - with pytest.raises(FileNotExistsError, match="One or more files not found"): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, - session=session, + session=sqlite_session, ) - def test_save_document_with_dataset_id_requires_notion_info_list_for_notion_import(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_requires_notion_info_list_for_notion_import( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.NOTION_IMPORT) + process_rule = _process_rule() + sqlite_session.add_all([dataset, process_rule]) + sqlite_session.commit() knowledge_config = KnowledgeConfig( indexing_technique="economy", data_source=DataSource(info_list=InfoList(data_source_type="notion_import")), @@ -1986,7 +2170,6 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch.object(DocumentService, "get_documents_position", return_value=1), ): mock_redis.lock.return_value = _make_lock_context() with pytest.raises(ValueError, match="No notion info list found"): @@ -1994,12 +2177,17 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: dataset, knowledge_config, account_context, - dataset_process_rule=SimpleNamespace(id="rule-1"), - session=MagicMock(), + dataset_process_rule=process_rule, + session=sqlite_session, ) - def test_save_document_with_dataset_id_requires_website_info_list_for_website_crawl(self, account_context): - dataset = _make_dataset() + def test_save_document_with_dataset_id_requires_website_info_list_for_website_crawl( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.WEBSITE_CRAWL) + process_rule = _process_rule() + sqlite_session.add_all([dataset, process_rule]) + sqlite_session.commit() knowledge_config = KnowledgeConfig( indexing_technique="economy", data_source=DataSource(info_list=InfoList(data_source_type="website_crawl")), @@ -2010,7 +2198,6 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: with ( patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), patch("services.dataset_service.redis_client") as mock_redis, - patch.object(DocumentService, "get_documents_position", return_value=1), ): mock_redis.lock.return_value = _make_lock_context() with pytest.raises(ValueError, match="No website info list found"): @@ -2018,6 +2205,6 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: dataset, knowledge_config, account_context, - dataset_process_rule=SimpleNamespace(id="rule-1"), - session=MagicMock(), + dataset_process_rule=process_rule, + session=sqlite_session, ) diff --git a/api/tests/unit_tests/services/test_feature_service_deployment_edition.py b/api/tests/unit_tests/services/test_feature_service_deployment_edition.py index 63ebb75c4ba..e0a9bd68685 100644 --- a/api/tests/unit_tests/services/test_feature_service_deployment_edition.py +++ b/api/tests/unit_tests/services/test_feature_service_deployment_edition.py @@ -43,3 +43,24 @@ def test_get_system_features_uses_configured_deployment_edition( fulfill_from_enterprise.assert_called_once_with(result) else: fulfill_from_enterprise.assert_not_called() + + +@pytest.mark.parametrize( + ("edition", "feature_enabled", "expected"), + [ + (DeploymentEdition.CLOUD, True, True), + (DeploymentEdition.CLOUD, False, False), + (DeploymentEdition.COMMUNITY, True, False), + (DeploymentEdition.ENTERPRISE, True, False), + ], +) +def test_trial_app_policy_is_cloud_only( + monkeypatch: pytest.MonkeyPatch, + edition: DeploymentEdition, + feature_enabled: bool, + expected: bool, +) -> None: + monkeypatch.setattr("services.feature_service.dify_config.DEPLOYMENT_EDITION", edition) + monkeypatch.setattr("services.feature_service.dify_config.ENABLE_TRIAL_APP", feature_enabled) + + assert FeatureService.is_trial_app_enabled() is expected diff --git a/api/tests/unit_tests/services/test_file_service.py b/api/tests/unit_tests/services/test_file_service.py index 0994abecf65..f345a2ec012 100644 --- a/api/tests/unit_tests/services/test_file_service.py +++ b/api/tests/unit_tests/services/test_file_service.py @@ -11,6 +11,7 @@ from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import NotFound from configs import dify_config +from enums import DeploymentEdition from extensions.storage.storage_type import StorageType from models.base import TypeBase from models.enums import CreatorUserRole @@ -296,6 +297,40 @@ class TestFileService: with pytest.raises(NotFound, match="File not found"): file_service.get_file_presigned_url(file_id="file_id", tenant_id="tenant_id") + def test_get_icon_url_uses_direct_storage_url_for_cloud_s3(self, file_service: FileService): + with ( + patch.object(dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.CLOUD), + patch.object(dify_config, "STORAGE_TYPE", StorageType.S3), + patch.object(file_service, "get_file_presigned_url", return_value="direct-url") as get_presigned_url, + ): + result = file_service.get_icon_url("file_id", "tenant_id") + + assert result == "direct-url" + get_presigned_url.assert_called_once_with(file_id="file_id", tenant_id="tenant_id") + + @pytest.mark.parametrize( + ("deployment_edition", "storage_type"), + [ + (DeploymentEdition.COMMUNITY, StorageType.S3), + (DeploymentEdition.CLOUD, StorageType.LOCAL), + ], + ) + def test_get_icon_url_uses_preview_url_outside_cloud_s3( + self, + file_service: FileService, + deployment_edition: DeploymentEdition, + storage_type: StorageType, + ): + with ( + patch.object(dify_config, "DEPLOYMENT_EDITION", deployment_edition), + patch.object(dify_config, "STORAGE_TYPE", storage_type), + patch("services.file_service.file_helpers.get_signed_file_url", return_value="preview-url") as get_url, + ): + result = file_service.get_icon_url("file_id", "tenant_id") + + assert result == "preview-url" + get_url.assert_called_once_with(upload_file_id="file_id") + def test_upload_text_success(self, file_service: FileService, db_session: Session): # Setup text = "sample text" diff --git a/api/tests/unit_tests/services/test_message_service.py b/api/tests/unit_tests/services/test_message_service.py index e0a8402b84e..e6e5c76a9ca 100644 --- a/api/tests/unit_tests/services/test_message_service.py +++ b/api/tests/unit_tests/services/test_message_service.py @@ -38,6 +38,7 @@ from models.model import ( Message, MessageFeedback, ) +from models.workflow import Workflow, WorkflowType from repositories.sqlalchemy_execution_extra_content_repository import SQLAlchemyExecutionExtraContentRepository from services.errors.message import ( FirstMessageNotExistsError, @@ -111,6 +112,19 @@ class MessageServiceTestDataFactory: account.id = user_id return account + @staticmethod + def create_workflow(*, features: dict[str, object] | None = None) -> Workflow: + return Workflow( + id="workflow-123", + tenant_id="tenant-123", + app_id="app-123", + type=WorkflowType.CHAT, + version="1", + graph="{}", + _features=json.dumps(features or {}), + created_by="account-123", + ) + @staticmethod def create_conversation( conversation_id: str = "conv-001", @@ -712,8 +726,7 @@ class TestMessageServiceSuggestedQuestions: monkeypatch: pytest.MonkeyPatch, conversation: Conversation, ) -> tuple[MagicMock, MagicMock, MagicMock]: - message = MagicMock() - message.conversation_id = conversation.id + message = MessageServiceTestDataFactory.create_message(message_id="msg-123", conversation_id=conversation.id) monkeypatch.setattr(service_module.MessageService, "get_message", MagicMock(return_value=message)) monkeypatch.setattr( service_module.ConversationService, "get_conversation", MagicMock(return_value=conversation) @@ -747,8 +760,7 @@ class TestMessageServiceSuggestedQuestions: ) -> None: conversation = factory.create_conversation() _, _, llm_generator = self._chat_boundaries(monkeypatch, conversation) - workflow = MagicMock() - workflow.features_dict = {"suggested_questions_after_answer": {"enabled": True}} + workflow = factory.create_workflow(features={"suggested_questions_after_answer": {"enabled": True}}) workflow_service = MagicMock() workflow_service.return_value.get_published_workflow.return_value = workflow monkeypatch.setattr(service_module, "WorkflowService", workflow_service) @@ -1103,7 +1115,7 @@ class TestMessageServiceSuggestedQuestions: ) -> None: conversation = factory.create_conversation() self._chat_boundaries(monkeypatch, conversation) - workflow = MagicMock() + workflow = factory.create_workflow() workflow_service = MagicMock() workflow_service.return_value.get_published_workflow.return_value = workflow monkeypatch.setattr(service_module, "WorkflowService", workflow_service) diff --git a/api/tests/unit_tests/services/test_recommended_app_catalog_gateway.py b/api/tests/unit_tests/services/test_recommended_app_catalog_gateway.py new file mode 100644 index 00000000000..bcd48c1b2d3 --- /dev/null +++ b/api/tests/unit_tests/services/test_recommended_app_catalog_gateway.py @@ -0,0 +1,661 @@ +import json +from collections.abc import Iterator +from pathlib import Path +from unittest.mock import MagicMock, call, patch + +import pytest +import yaml + +from services import recommended_app_catalog_gateway as gateway_module +from services.recommended_app_catalog_gateway import ( + BuiltinRecommendedAppCatalogGateway, + RecommendedAppCatalogRouter, + RemoteRecommendedAppCatalogGateway, +) +from services.recommended_app_query_service import ( + RecommendedAppCatalogPage, + RecommendedAppDetailRecord, + RecommendedAppInfoRecord, + RecommendedAppRecord, +) + + +def _page_payload(*app_ids: str, learn_dify_ids: frozenset[str] = frozenset()) -> dict[str, object]: + app_ids = app_ids or ("app-1",) + return { + "recommended_apps": [ + { + "app": { + "id": app_id, + "name": "App", + "mode": "chat", + "icon": "icon.png", + "icon_type": "image", + "icon_background": "#fff", + }, + "app_id": app_id, + "description": "description", + "copyright": None, + "privacy_policy": None, + "categories": ["Workflow"], + "position": 1, + "is_listed": True, + **({"is_learn_dify": True} if app_id in learn_dify_ids else {}), + } + for app_id in app_ids + ], + "categories": ["Workflow"], + } + + +def _detail_payload() -> dict[str, object]: + return { + "id": "app-1", + "name": "App", + "icon": None, + "icon_background": None, + "mode": "chat", + "export_data": "{}", + } + + +def _expected_page(*, categories: tuple[str, ...] = ("Workflow",)) -> RecommendedAppCatalogPage: + return RecommendedAppCatalogPage( + recommended_apps=( + RecommendedAppRecord( + app=RecommendedAppInfoRecord( + id="app-1", + name="App", + mode="chat", + icon="icon.png", + icon_type="image", + icon_background="#fff", + ), + app_id="app-1", + description="description", + copyright=None, + privacy_policy=None, + custom_disclaimer=None, + categories=("Workflow",), + position=1, + is_listed=True, + ), + ), + categories=categories, + ) + + +class TestBuiltinRecommendedAppCatalogGateway: + def test_maps_bundled_catalog(self) -> None: + gateway = BuiltinRecommendedAppCatalogGateway() + page = gateway.list_recommended("en-US") + learn_dify = gateway.list_learn_dify("ja-JP") + detail = gateway.get_detail(page.recommended_apps[0].app_id) + + assert page.recommended_apps + assert [app.app_id for app in learn_dify.recommended_apps] == [ + "f00c4531-6551-45ee-808f-1d7903099515", + "d9f6b733-e35d-4a40-9f38-ca7bbfa009f7", + "e9870913-dd01-4710-9f06-15d4180ca1ce", + ] + assert all(gateway.get_detail(app.app_id) is not None for app in learn_dify.recommended_apps) + assert detail is not None + + def test_bundled_workflow_templates_have_unique_end_output_variables(self) -> None: + data_path = Path(gateway_module.__file__).resolve().parents[1] / "constants" / "recommended_apps.json" + data = json.loads(data_path.read_text(encoding="utf-8")) + + offenders: dict[str, list[str]] = {} + for app_id, detail in data.get("app_details", {}).items(): + export_data = detail.get("export_data") + if not export_data: + continue + dsl = yaml.safe_load(export_data) + nodes = (dsl or {}).get("workflow", {}).get("graph", {}).get("nodes", []) + output_names = [ + output.get("variable") + for node in nodes + if node.get("data", {}).get("type") == "end" + for output in (node.get("data", {}).get("outputs") or []) + ] + duplicates = sorted({name for name in output_names if output_names.count(name) > 1}) + if duplicates: + offenders[detail.get("name", app_id).strip()] = duplicates + + assert offenders == {}, f"templates with duplicate End output variable names: {offenders}" + + def test_maps_builtin_payload(self) -> None: + gateway = BuiltinRecommendedAppCatalogGateway() + gateway._data = { + "recommended_apps": {"en-US": _page_payload("app-1", "app-2", learn_dify_ids=frozenset({"app-1"}))}, + "app_details": {"app-1": _detail_payload()}, + } + + assert [app.app_id for app in gateway.list_recommended("en-US").recommended_apps] == ["app-1", "app-2"] + assert gateway.list_learn_dify("en-US") == RecommendedAppCatalogPage( + recommended_apps=_expected_page().recommended_apps, + categories=(), + ) + assert gateway.get_detail("app-1") == RecommendedAppDetailRecord( + id="app-1", + name="App", + icon=None, + icon_background=None, + mode="chat", + export_data="{}", + ) + + def test_membership_uses_raw_non_none_detail(self) -> None: + gateway = BuiltinRecommendedAppCatalogGateway() + gateway._data = {"app_details": {"malformed": object()}} + + assert gateway.contains("malformed") is True + assert gateway.contains("missing") is False + with pytest.raises(TypeError, match="recommended app detail must be a mapping"): + gateway.get_detail("malformed") + + def test_missing_language_returns_empty_page(self) -> None: + gateway = BuiltinRecommendedAppCatalogGateway() + gateway._data = {"recommended_apps": {}} + + assert gateway.list_recommended("fr-FR") == RecommendedAppCatalogPage( + recommended_apps=(), + categories=(), + ) + assert gateway.list_learn_dify("fr-FR") == RecommendedAppCatalogPage( + recommended_apps=(), + categories=(), + ) + + def test_nonempty_page_requires_categories(self) -> None: + page = _page_payload() + del page["categories"] + gateway = BuiltinRecommendedAppCatalogGateway() + gateway._data = {"recommended_apps": {"en-US": page}} + + with pytest.raises(KeyError, match="categories"): + gateway.list_recommended("en-US") + + @pytest.mark.parametrize("categories", ["Agent", b"Agent", ["Agent", 1]]) + def test_rejects_malformed_page_categories(self, categories: object) -> None: + page = _page_payload() + page["categories"] = categories + gateway = BuiltinRecommendedAppCatalogGateway() + gateway._data = {"recommended_apps": {"en-US": page}} + + with pytest.raises(TypeError, match="categories must"): + gateway.list_recommended("en-US") + + def test_rejects_malformed_app_categories(self) -> None: + gateway = BuiltinRecommendedAppCatalogGateway() + gateway._data = { + "recommended_apps": { + "en-US": { + "recommended_apps": [{"app": None, "app_id": "app-1", "categories": "Agent"}], + "categories": ["Agent"], + } + } + } + + with pytest.raises(TypeError, match="categories must"): + gateway.list_recommended("en-US") + + def test_rejects_non_string_detail_mode(self) -> None: + detail = _detail_payload() + detail["mode"] = object() + gateway = BuiltinRecommendedAppCatalogGateway() + gateway._data = {"app_details": {"app-1": detail}} + + with pytest.raises(TypeError, match="mode must be a string"): + gateway.get_detail("app-1") + + def test_reads_builtin_file_once_per_gateway(self) -> None: + gateway = BuiltinRecommendedAppCatalogGateway() + payload = json.dumps({"recommended_apps": {"en-US": _page_payload()}}) + + with patch.object(gateway_module.Path, "read_text", return_value=payload) as read_text: + gateway.list_recommended("en-US") + gateway.list_recommended("en-US") + + read_text.assert_called_once_with(encoding="utf-8") + + +class TestRemoteRecommendedAppCatalogGateway: + @pytest.fixture(autouse=True) + def _use_remote_mode(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + gateway_module.clear_remote_fetch_cache() + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "remote") + yield + gateway_module.clear_remote_fetch_cache() + + def test_maps_remote_pages_without_reordering(self, monkeypatch: pytest.MonkeyPatch) -> None: + gateway = RemoteRecommendedAppCatalogGateway() + payload = _page_payload("app-2", "app-1") + payload["categories"] = ["Writing", "Agent"] + monkeypatch.setattr(gateway, "_fetch_page", MagicMock(return_value=payload)) + monkeypatch.setattr(gateway, "_fetch_learn_dify_page", MagicMock(return_value=payload)) + + recommended = gateway.list_recommended("en-US") + learn_dify = gateway.list_learn_dify("en-US") + assert [app.app_id for app in recommended.recommended_apps] == ["app-2", "app-1"] + assert recommended.categories == ("Writing", "Agent") + assert [app.app_id for app in learn_dify.recommended_apps] == ["app-2", "app-1"] + assert learn_dify.categories == () + + def test_list_fetch_error_falls_back_through_builtin_en_us(self, monkeypatch: pytest.MonkeyPatch) -> None: + fallback = MagicMock() + empty_page = RecommendedAppCatalogPage(recommended_apps=(), categories=()) + fallback_page = _expected_page(categories=("builtin",)) + fallback.list_recommended.side_effect = [empty_page, fallback_page] + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + + monkeypatch.setattr(remote, "_fetch_page", MagicMock(side_effect=ConnectionError("timeout"))) + + assert router.list_recommended("fr-FR") == fallback_page + assert fallback.list_recommended.call_args_list == [call("fr-FR"), call("en-US")] + + def test_json_decode_error_falls_back_to_builtin(self, monkeypatch: pytest.MonkeyPatch) -> None: + fallback = MagicMock() + expected_page = _expected_page() + fallback.list_recommended.return_value = expected_page + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + response = MagicMock(status_code=200) + response.json.side_effect = ValueError("invalid JSON") + monkeypatch.setattr(gateway_module.httpx, "get", MagicMock(return_value=response)) + + assert router.list_recommended("en-US") == expected_page + fallback.list_recommended.assert_called_once_with("en-US") + + def test_payload_mapping_error_does_not_fall_back(self, monkeypatch: pytest.MonkeyPatch) -> None: + fallback = MagicMock() + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + monkeypatch.setattr(remote, "_fetch_page", MagicMock(return_value=object())) + + with pytest.raises(TypeError, match="recommended app page must be a mapping"): + router.list_recommended("en-US") + fallback.list_recommended.assert_not_called() + + def test_learn_dify_fetch_error_falls_back_to_builtin(self, monkeypatch: pytest.MonkeyPatch) -> None: + builtin = MagicMock() + database = MagicMock() + page = RecommendedAppCatalogPage(recommended_apps=(), categories=()) + builtin.list_learn_dify.return_value = page + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=database, + builtin=builtin, + ) + + monkeypatch.setattr(remote, "_fetch_learn_dify_page", MagicMock(side_effect=ConnectionError("timeout"))) + + assert router.list_learn_dify("ja-JP") == page + builtin.list_learn_dify.assert_called_once_with("ja-JP") + database.list_learn_dify.assert_not_called() + + def test_empty_remote_learn_dify_page_does_not_fall_back(self, monkeypatch: pytest.MonkeyPatch) -> None: + database = MagicMock() + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=database, + builtin=MagicMock(), + ) + monkeypatch.setattr( + remote, + "_fetch_learn_dify_page", + MagicMock(return_value={"recommended_apps": [], "categories": []}), + ) + + assert router.list_learn_dify("en-US") == RecommendedAppCatalogPage( + recommended_apps=(), + categories=(), + ) + database.list_learn_dify.assert_not_called() + + @pytest.mark.parametrize("status_code", [404, 500]) + def test_detail_non_200_returns_none_without_builtin_fallback( + self, + monkeypatch: pytest.MonkeyPatch, + status_code: int, + ) -> None: + fallback = MagicMock() + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + response = MagicMock(status_code=status_code) + monkeypatch.setattr("services.recommended_app_catalog_gateway.httpx.get", MagicMock(return_value=response)) + + assert router.get_detail("missing") is None + fallback.get_detail.assert_not_called() + + def test_detail_fetch_error_falls_back_to_builtin(self, monkeypatch: pytest.MonkeyPatch) -> None: + fallback = MagicMock() + fallback_detail = RecommendedAppDetailRecord( + id="fallback", + name="Fallback", + icon=None, + icon_background=None, + mode="chat", + export_data="{}", + ) + fallback.get_detail.return_value = fallback_detail + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + + monkeypatch.setattr(remote, "_fetch_detail", MagicMock(side_effect=ConnectionError("timeout"))) + + assert router.get_detail("app-1") == fallback_detail + fallback.get_detail.assert_called_once_with("app-1") + + def test_detail_mapping_error_does_not_fall_back(self, monkeypatch: pytest.MonkeyPatch) -> None: + fallback = MagicMock() + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + monkeypatch.setattr(remote, "_fetch_detail", MagicMock(return_value=object())) + + with pytest.raises(TypeError, match="recommended app detail must be a mapping"): + router.get_detail("app-1") + fallback.get_detail.assert_not_called() + + def test_learn_dify_mapping_error_does_not_fall_back(self, monkeypatch: pytest.MonkeyPatch) -> None: + database = MagicMock() + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=database, + builtin=MagicMock(), + ) + monkeypatch.setattr(remote, "_fetch_learn_dify_page", MagicMock(return_value=object())) + + with pytest.raises(TypeError, match="Learn Dify app page must be a mapping"): + router.list_learn_dify("en-US") + database.list_learn_dify.assert_not_called() + + def test_membership_accepts_raw_non_none_payload(self, monkeypatch: pytest.MonkeyPatch) -> None: + gateway = RemoteRecommendedAppCatalogGateway() + monkeypatch.setattr(gateway, "_fetch_detail", MagicMock(return_value=object())) + + assert gateway.contains("app-1") is True + + @pytest.mark.parametrize("status_code", [404, 500]) + def test_membership_non_200_does_not_fall_back( + self, + monkeypatch: pytest.MonkeyPatch, + status_code: int, + ) -> None: + fallback = MagicMock() + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + response = MagicMock(status_code=status_code) + monkeypatch.setattr(gateway_module.httpx, "get", MagicMock(return_value=response)) + + assert router.contains("missing") is False + fallback.contains.assert_not_called() + + def test_membership_fetch_error_falls_back_to_builtin(self, monkeypatch: pytest.MonkeyPatch) -> None: + fallback = MagicMock() + fallback.contains.return_value = True + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=fallback, + ) + monkeypatch.setattr(remote, "_fetch_detail", MagicMock(side_effect=ConnectionError("timeout"))) + + assert router.contains("app-1") is True + fallback.contains.assert_called_once_with("app-1") + + def test_remote_request_uses_configured_origin_and_timeouts( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + response = MagicMock(status_code=200) + response.json.return_value = _detail_payload() + http_get = MagicMock(return_value=response) + monkeypatch.setattr(gateway_module.httpx, "get", http_get) + monkeypatch.setattr( + gateway_module.dify_config, + "HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN", + "https://catalog.example.com", + ) + monkeypatch.setattr(gateway_module.dify_config, "CONSOLE_WEB_URL", "https://console.example.com") + gateway = RemoteRecommendedAppCatalogGateway() + gateway.get_detail("app-1") + + http_get.assert_called_once() + call = http_get.call_args + assert call.args == ("https://catalog.example.com/apps/app-1",) + assert call.kwargs["headers"] == {"Origin": "https://console.example.com"} + assert call.kwargs["timeout"].connect == 3.0 + assert call.kwargs["timeout"].read == 10.0 + + def test_remote_request_uses_cache(self, monkeypatch: pytest.MonkeyPatch) -> None: + response = MagicMock(status_code=200) + response.json.return_value = _page_payload() + http_get = MagicMock(return_value=response) + monkeypatch.setattr(gateway_module.httpx, "get", http_get) + monkeypatch.setattr( + gateway_module.dify_config, + "HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN", + "https://catalog.example.com", + ) + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL", 600) + gateway = RemoteRecommendedAppCatalogGateway() + + assert gateway.list_recommended("en-US") == _expected_page() + assert gateway.list_recommended("en-US") == _expected_page() + http_get.assert_called_once() + + def test_remote_request_does_not_cache_failed_responses(self, monkeypatch: pytest.MonkeyPatch) -> None: + response = MagicMock(status_code=500) + http_get = MagicMock(return_value=response) + monkeypatch.setattr(gateway_module.httpx, "get", http_get) + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL", 600) + expected_page = _expected_page() + fallback = MagicMock() + fallback.list_recommended.return_value = expected_page + router = RecommendedAppCatalogRouter( + remote=RemoteRecommendedAppCatalogGateway(), + database=MagicMock(), + builtin=fallback, + ) + + assert router.list_recommended("en-US") == expected_page + assert router.list_recommended("en-US") == expected_page + assert http_get.call_count == 2 + + def test_remote_request_skips_cache_when_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None: + response = MagicMock(status_code=200) + response.json.return_value = _page_payload() + http_get = MagicMock(return_value=response) + monkeypatch.setattr(gateway_module.httpx, "get", http_get) + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL", 0) + gateway = RemoteRecommendedAppCatalogGateway() + + gateway.list_recommended("en-US") + gateway.list_recommended("en-US") + assert http_get.call_count == 2 + + def test_remote_request_cache_isolated_by_configured_origin( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + response = MagicMock(status_code=200) + response.json.return_value = _page_payload() + http_get = MagicMock(return_value=response) + monkeypatch.setattr(gateway_module.httpx, "get", http_get) + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_CACHE_TTL", 600) + gateway = RemoteRecommendedAppCatalogGateway() + monkeypatch.setattr(gateway_module.dify_config, "CONSOLE_WEB_URL", "https://cloud-a.example.com") + gateway.list_recommended("en-US") + monkeypatch.setattr(gateway_module.dify_config, "CONSOLE_WEB_URL", "https://cloud-b.example.com") + gateway.list_recommended("en-US") + + assert http_get.call_count == 2 + + @pytest.mark.parametrize( + ("console_web_url", "expected_headers"), + [ + ("saas.dify.dev", {"Origin": "saas.dify.dev"}), + ("http://localhost:3000/console", {"Origin": "http://localhost:3000/console"}), + ("", {}), + ], + ) + def test_remote_request_uses_console_web_url( + self, + monkeypatch: pytest.MonkeyPatch, + console_web_url: str, + expected_headers: dict[str, str], + ) -> None: + response = MagicMock(status_code=200) + response.json.return_value = _detail_payload() + http_get = MagicMock(return_value=response) + monkeypatch.setattr(gateway_module.httpx, "get", http_get) + monkeypatch.setattr(gateway_module.dify_config, "CONSOLE_WEB_URL", console_web_url) + gateway = RemoteRecommendedAppCatalogGateway() + gateway.get_detail("app-1") + + assert http_get.call_args.kwargs["headers"] == expected_headers + + @pytest.mark.parametrize( + ("operation", "expected_url"), + [ + ("recommended", "https://catalog.example.com/apps?language=ja-JP"), + ("learn_dify", "https://catalog.example.com/apps/learn-dify?language=ja-JP"), + ], + ) + def test_remote_list_non_200_uses_expected_fallback( + self, + monkeypatch: pytest.MonkeyPatch, + operation: str, + expected_url: str, + ) -> None: + response = MagicMock(status_code=500) + http_get = MagicMock(return_value=response) + monkeypatch.setattr(gateway_module.httpx, "get", http_get) + monkeypatch.setattr( + gateway_module.dify_config, + "HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN", + "https://catalog.example.com", + ) + fallback = MagicMock() + database = MagicMock() + expected_page = _expected_page() + fallback.list_recommended.return_value = expected_page + fallback.list_learn_dify.return_value = expected_page + remote = RemoteRecommendedAppCatalogGateway() + router = RecommendedAppCatalogRouter( + remote=remote, + database=database, + builtin=fallback, + ) + + result = router.list_recommended("ja-JP") if operation == "recommended" else router.list_learn_dify("ja-JP") + fallback_call = fallback.list_recommended if operation == "recommended" else fallback.list_learn_dify + + assert result == expected_page + assert http_get.call_args.args == (expected_url,) + fallback_call.assert_called_once_with("ja-JP") + database.list_learn_dify.assert_not_called() + + +class TestRecommendedAppCatalogRouter: + def test_empty_page_falls_back_to_builtin_en_us(self, monkeypatch: pytest.MonkeyPatch) -> None: + remote = MagicMock() + builtin = MagicMock() + remote.list_recommended.return_value = RecommendedAppCatalogPage(recommended_apps=(), categories=()) + expected_page = _expected_page(categories=("builtin",)) + builtin.list_recommended.return_value = expected_page + gateway = RecommendedAppCatalogRouter( + remote=remote, + database=MagicMock(), + builtin=builtin, + ) + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "remote") + + assert gateway.list_recommended("ja-JP") == expected_page + remote.list_recommended.assert_called_once_with("ja-JP") + builtin.list_recommended.assert_called_once_with("en-US") + + def test_resolves_mode_for_every_operation(self, monkeypatch: pytest.MonkeyPatch) -> None: + remote = MagicMock() + database = MagicMock() + builtin = MagicMock() + gateway = RecommendedAppCatalogRouter( + remote=remote, + database=database, + builtin=builtin, + ) + + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "remote") + gateway.list_recommended("en-US") + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "db") + gateway.list_learn_dify("en-US") + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "builtin") + gateway.get_detail("app-1") + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "remote") + gateway.contains("app-1") + + remote.list_recommended.assert_called_once_with("en-US") + database.list_learn_dify.assert_called_once_with("en-US") + builtin.get_detail.assert_called_once_with("app-1") + remote.contains.assert_called_once_with("app-1") + + def test_builtin_mode_reads_builtin_learn_dify(self, monkeypatch: pytest.MonkeyPatch) -> None: + builtin = MagicMock() + database = MagicMock() + expected_page = RecommendedAppCatalogPage(recommended_apps=(), categories=()) + builtin.list_learn_dify.return_value = expected_page + gateway = RecommendedAppCatalogRouter( + remote=MagicMock(), + database=database, + builtin=builtin, + ) + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "builtin") + + assert gateway.list_learn_dify("en-US") == expected_page + builtin.list_learn_dify.assert_called_once_with("en-US") + database.list_learn_dify.assert_not_called() + + def test_rejects_invalid_mode(self, monkeypatch: pytest.MonkeyPatch) -> None: + gateway = RecommendedAppCatalogRouter( + remote=MagicMock(), + database=MagicMock(), + builtin=MagicMock(), + ) + monkeypatch.setattr(gateway_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "invalid") + + with pytest.raises(ValueError, match="invalid fetch recommended apps mode: invalid"): + gateway.list_recommended("en-US") diff --git a/api/tests/unit_tests/services/test_recommended_app_query_service.py b/api/tests/unit_tests/services/test_recommended_app_query_service.py new file mode 100644 index 00000000000..09d661b6754 --- /dev/null +++ b/api/tests/unit_tests/services/test_recommended_app_query_service.py @@ -0,0 +1,201 @@ +from unittest.mock import MagicMock + +import pytest + +from constants.languages import languages +from services.recommended_app_query_service import ( + RecommendedAppCatalogPage, + RecommendedAppDetailRecord, + RecommendedAppInfoRecord, + RecommendedAppNotFoundError, + RecommendedAppQueryService, + RecommendedAppRecord, +) + + +def _app(app_id: str) -> RecommendedAppRecord: + return RecommendedAppRecord( + app=RecommendedAppInfoRecord( + id=app_id, + name=f"App {app_id}", + mode="chat", + icon="icon.png", + icon_type="image", + icon_background="#fff", + ), + app_id=app_id, + description="description", + copyright=None, + privacy_policy=None, + custom_disclaimer=None, + categories=("Workflow",), + position=1, + is_listed=True, + ) + + +def _page(*app_ids: str, categories: tuple[str, ...] = ("Workflow",)) -> RecommendedAppCatalogPage: + return RecommendedAppCatalogPage( + recommended_apps=tuple(_app(app_id) for app_id in app_ids), + categories=categories, + ) + + +def _service( + *, + catalog: MagicMock, + trial_apps: MagicMock | None = None, + trial_enabled: bool = False, +) -> tuple[RecommendedAppQueryService, MagicMock]: + trial_apps = trial_apps or MagicMock() + return ( + RecommendedAppQueryService( + catalog=catalog, + trial_apps=trial_apps, + trial_enabled=trial_enabled, + ), + trial_apps, + ) + + +def test_is_previewable_accepts_trial_registration_without_querying_catalog() -> None: + catalog = MagicMock() + trial_apps = MagicMock() + trial_apps.existing_ids.return_value = frozenset({"app-1"}) + service, _ = _service(catalog=catalog, trial_apps=trial_apps) + + assert service.is_previewable("app-1") is True + trial_apps.existing_ids.assert_called_once_with(("app-1",)) + catalog.contains.assert_not_called() + + +@pytest.mark.parametrize("expected", [True, False]) +def test_is_previewable_falls_back_to_catalog(expected: bool) -> None: + catalog = MagicMock() + catalog.contains.return_value = expected + trial_apps = MagicMock() + trial_apps.existing_ids.return_value = frozenset() + service, _ = _service(catalog=catalog, trial_apps=trial_apps) + + assert service.is_previewable("app-1") is expected + trial_apps.existing_ids.assert_called_once_with(("app-1",)) + catalog.contains.assert_called_once_with("app-1") + + +@pytest.mark.parametrize( + ("requested_language", "interface_language", "expected"), + [ + ("en-US", "fr-FR", "en-US"), + ("invalid", "fr-FR", "fr-FR"), + (None, "custom-language", "custom-language"), + (None, None, languages[0]), + ], +) +def test_list_recommended_resolves_language( + requested_language: str | None, + interface_language: str | None, + expected: str, +) -> None: + catalog = MagicMock() + catalog.list_recommended.return_value = _page("app-1") + service, _ = _service(catalog=catalog) + + service.list_recommended( + requested_language=requested_language, + interface_language=interface_language, + ) + + catalog.list_recommended.assert_called_once_with(expected) + + +def test_list_recommended_disables_upstream_trial_without_querying_trial_apps() -> None: + catalog = MagicMock() + catalog.list_recommended.return_value = _page("app-1") + service, trial_apps = _service(catalog=catalog) + + result = service.list_recommended(requested_language="en-US", interface_language=None) + + assert result.recommended_apps[0].can_trial is False + trial_apps.existing_ids.assert_not_called() + + +def test_list_recommended_enriches_trial_status_in_one_bulk_query() -> None: + catalog = MagicMock() + catalog.list_recommended.return_value = _page("app-1", "app-2") + trial_apps = MagicMock() + trial_apps.existing_ids.return_value = frozenset({"app-1"}) + service, _ = _service( + catalog=catalog, + trial_apps=trial_apps, + trial_enabled=True, + ) + + result = service.list_recommended(requested_language="en-US", interface_language=None) + + assert [app.can_trial for app in result.recommended_apps] == [True, False] + trial_apps.existing_ids.assert_called_once_with(["app-1", "app-2"]) + + +def test_list_learn_dify_does_not_return_categories() -> None: + catalog = MagicMock() + catalog.list_learn_dify.return_value = _page(categories=("ignored",)) + service, _ = _service(catalog=catalog) + + result = service.list_learn_dify(requested_language="invalid", interface_language="fr-FR") + + catalog.list_learn_dify.assert_called_once_with("fr-FR") + assert result.recommended_apps == () + assert not hasattr(result, "categories") + + +def test_get_detail_raises_not_found_without_querying_trial_apps() -> None: + catalog = MagicMock() + catalog.get_detail.return_value = None + service, trial_apps = _service(catalog=catalog, trial_enabled=True) + + with pytest.raises(RecommendedAppNotFoundError): + service.get_detail("missing") + trial_apps.existing_ids.assert_not_called() + + +def test_get_detail_does_not_query_trial_apps_when_disabled() -> None: + catalog = MagicMock() + catalog.get_detail.return_value = RecommendedAppDetailRecord( + id="catalog-app-id", + name="App", + icon=None, + icon_background=None, + mode="chat", + export_data="{}", + ) + service, trial_apps = _service(catalog=catalog) + + result = service.get_detail("route-app-id") + + assert result.can_trial is False + trial_apps.existing_ids.assert_not_called() + + +@pytest.mark.parametrize(("existing_ids", "expected"), [(frozenset({"catalog-app-id"}), True), (frozenset(), False)]) +def test_get_detail_uses_catalog_result_id_for_trial_status(existing_ids: frozenset[str], expected: bool) -> None: + catalog = MagicMock() + catalog.get_detail.return_value = RecommendedAppDetailRecord( + id="catalog-app-id", + name="App", + icon=None, + icon_background=None, + mode="chat", + export_data="{}", + ) + trial_apps = MagicMock() + trial_apps.existing_ids.return_value = existing_ids + service, _ = _service( + catalog=catalog, + trial_apps=trial_apps, + trial_enabled=True, + ) + + result = service.get_detail("route-app-id") + + assert result.can_trial is expected + trial_apps.existing_ids.assert_called_once_with(("catalog-app-id",)) diff --git a/api/tests/unit_tests/services/test_recommended_app_service.py b/api/tests/unit_tests/services/test_recommended_app_service.py deleted file mode 100644 index c4a039e1bc8..00000000000 --- a/api/tests/unit_tests/services/test_recommended_app_service.py +++ /dev/null @@ -1,562 +0,0 @@ -"""Unit tests for recommended app orchestration and SQLite-backed trial state.""" - -from __future__ import annotations - -import uuid -from collections.abc import Callable -from typing import TypedDict, Unpack, cast -from unittest.mock import MagicMock, patch - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from enums import DeploymentEdition -from models.model import AccountTrialAppRecord, App, AppMode, TrialApp -from services import recommended_app_service as service_module -from services.recommended_app_service import RecommendedAppService - -pytestmark = pytest.mark.parametrize( - "sqlite_session", - [(TrialApp, AccountTrialAppRecord, App)], - indirect=True, -) - - -@pytest.fixture(autouse=True) -def _recommended_app_config(config_overrides: Callable[..., None]) -> None: - config_overrides( - HOSTED_FETCH_APP_TEMPLATES_MODE="remote", - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - ENABLE_TRIAL_APP=True, - ) - - -class RecommendedAppPayload(TypedDict, total=False): - id: str - app_id: str - name: str - description: str - category: str - icon: str - model_config: object - workflows: list[str] - tools: list[str] - can_trial: bool - - -class AppsResponse(TypedDict): - recommended_apps: list[RecommendedAppPayload] | None - categories: list[str] - - -class AppDetailKwargs(TypedDict, total=False): - category: str - icon: str - model_config: object - workflows: list[str] - tools: list[str] - - -@pytest.mark.parametrize( - ("edition", "feature_enabled", "expected"), - [ - (DeploymentEdition.CLOUD, True, True), - (DeploymentEdition.CLOUD, False, False), - (DeploymentEdition.COMMUNITY, True, False), - (DeploymentEdition.ENTERPRISE, True, False), - ], -) -def test_trial_app_policy_is_cloud_only( - config_overrides: Callable[..., None], - sqlite_session: Session, - edition: DeploymentEdition, - feature_enabled: bool, - expected: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition, ENABLE_TRIAL_APP=feature_enabled) - - assert RecommendedAppService.is_trial_app_enabled() is expected - - -# ── Helpers ──────────────────────────────────────────────────────────── - - -def _apps_response( - recommended_apps: list[RecommendedAppPayload] | None = None, - categories: list[str] | None = None, -) -> AppsResponse: - if recommended_apps is None: - recommended_apps = [ - {"app_id": "app-1", "name": "Test App 1", "description": "d1", "category": "productivity"}, - {"app_id": "app-2", "name": "Test App 2", "description": "d2", "category": "communication"}, - ] - if categories is None: - categories = ["productivity", "communication", "utilities"] - return {"recommended_apps": recommended_apps, "categories": categories} - - -def _app_detail( - app_id: str = "app-123", - name: str = "Test App", - description: str = "Test description", - **kwargs: Unpack[AppDetailKwargs], -) -> RecommendedAppPayload: - detail = RecommendedAppPayload( - id=app_id, - name=name, - description=description, - category=kwargs.get("category", "productivity"), - icon=kwargs.get("icon", "🚀"), - model_config=kwargs.get("model_config", {}), - ) - detail.update(**kwargs) - return detail - - -def _mock_factory_for_apps( - monkeypatch: pytest.MonkeyPatch, - *, - mode: str, - result: AppsResponse, - fallback_result: AppsResponse | None = None, -) -> tuple[MagicMock, MagicMock]: - retrieval_instance = MagicMock() - retrieval_instance.get_recommended_apps_and_categories.return_value = result - retrieval_factory = MagicMock(return_value=retrieval_instance) - monkeypatch.setattr(service_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", mode, raising=False) - monkeypatch.setattr( - service_module.RecommendAppRetrievalFactory, - "get_recommend_app_factory", - MagicMock(return_value=retrieval_factory), - ) - builtin_instance = MagicMock() - if fallback_result is not None: - builtin_instance.fetch_recommended_apps_from_builtin.return_value = fallback_result - monkeypatch.setattr( - service_module.RecommendAppRetrievalFactory, - "get_buildin_recommend_app_retrieval", - MagicMock(return_value=builtin_instance), - ) - return retrieval_instance, builtin_instance - - -def _mock_factory_for_app_detail( - monkeypatch: pytest.MonkeyPatch, - *, - result: RecommendedAppPayload | None, -) -> MagicMock: - retrieval_instance = MagicMock() - retrieval_instance.get_recommend_app_detail.return_value = result - retrieval_factory = MagicMock(return_value=retrieval_instance) - monkeypatch.setattr(service_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "remote", raising=False) - monkeypatch.setattr( - service_module.RecommendAppRetrievalFactory, - "get_recommend_app_factory", - MagicMock(return_value=retrieval_factory), - ) - return retrieval_instance - - -def _persist_app(session: Session, *, name: str) -> App: - app = App( - tenant_id=str(uuid.uuid4()), - name=name, - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - ) - app.id = str(uuid.uuid4()) - session.add(app) - session.commit() - return app - - -# ── Pure logic tests: get_recommended_apps_and_categories ────────────── - - -class TestRecommendedAppServiceGetApps: - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_success_with_apps(self, mock_factory_class: MagicMock, sqlite_session: Session) -> None: - expected = _apps_response() - - mock_instance = MagicMock() - mock_instance.get_recommended_apps_and_categories.return_value = expected - mock_factory = MagicMock(return_value=mock_instance) - mock_factory_class.get_recommend_app_factory.return_value = mock_factory - - result = RecommendedAppService.get_recommended_apps_and_categories("en-US", session=sqlite_session) - - assert result == expected - assert len(result["recommended_apps"]) == 2 - assert len(result["categories"]) == 3 - mock_factory_class.get_recommend_app_factory.assert_called_once_with("remote") - mock_instance.get_recommended_apps_and_categories.assert_called_once_with("en-US", session=sqlite_session) - - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_fallback_to_builtin_when_empty(self, mock_factory_class: MagicMock, sqlite_session: Session) -> None: - empty_response = AppsResponse(recommended_apps=[], categories=[]) - builtin_response = _apps_response( - recommended_apps=[{"app_id": "builtin-1", "name": "Builtin App", "category": "default"}] - ) - - mock_remote_instance = MagicMock() - mock_remote_instance.get_recommended_apps_and_categories.return_value = empty_response - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_remote_instance) - - mock_builtin_instance = MagicMock() - mock_builtin_instance.fetch_recommended_apps_from_builtin.return_value = builtin_response - mock_factory_class.get_buildin_recommend_app_retrieval.return_value = mock_builtin_instance - - result = RecommendedAppService.get_recommended_apps_and_categories("zh-CN", session=sqlite_session) - - assert result == builtin_response - assert result["recommended_apps"][0]["app_id"] == "builtin-1" - mock_builtin_instance.fetch_recommended_apps_from_builtin.assert_called_once_with("en-US") - - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_fallback_when_none_recommended_apps( - self, - mock_factory_class: MagicMock, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(HOSTED_FETCH_APP_TEMPLATES_MODE="db") - none_response = AppsResponse(recommended_apps=None, categories=["test"]) - builtin_response = _apps_response() - - mock_db_instance = MagicMock() - mock_db_instance.get_recommended_apps_and_categories.return_value = none_response - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_db_instance) - - mock_builtin_instance = MagicMock() - mock_builtin_instance.fetch_recommended_apps_from_builtin.return_value = builtin_response - mock_factory_class.get_buildin_recommend_app_retrieval.return_value = mock_builtin_instance - - result = RecommendedAppService.get_recommended_apps_and_categories("en-US", session=sqlite_session) - - assert result == builtin_response - mock_builtin_instance.fetch_recommended_apps_from_builtin.assert_called_once() - - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_different_languages( - self, - mock_factory_class: MagicMock, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(HOSTED_FETCH_APP_TEMPLATES_MODE="builtin") - - for language in ["en-US", "zh-CN", "ja-JP", "fr-FR"]: - lang_response = _apps_response( - recommended_apps=[{"app_id": f"app-{language}", "name": f"App {language}", "category": "test"}] - ) - mock_instance = MagicMock() - mock_instance.get_recommended_apps_and_categories.return_value = lang_response - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - - result = RecommendedAppService.get_recommended_apps_and_categories(language, session=sqlite_session) - - assert result["recommended_apps"][0]["app_id"] == f"app-{language}" - mock_instance.get_recommended_apps_and_categories.assert_called_with(language, session=sqlite_session) - - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_uses_correct_factory_mode( - self, - mock_factory_class: MagicMock, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - for mode in ["remote", "builtin", "db"]: - config_overrides(HOSTED_FETCH_APP_TEMPLATES_MODE=mode) - response = _apps_response() - mock_instance = MagicMock() - mock_instance.get_recommended_apps_and_categories.return_value = response - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - - RecommendedAppService.get_recommended_apps_and_categories("en-US", session=sqlite_session) - - mock_factory_class.get_recommend_app_factory.assert_called_with(mode) - - -# ── Database-backed tests: get_app ───────────────────────────────────── - - -class TestRecommendedAppServiceGetApp: - def test_returns_normal_recommended_app(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - app = _persist_app(sqlite_session, name="Recommended App") - - retrieval_instance = _mock_factory_for_app_detail( - monkeypatch, - result=RecommendedAppPayload(id=app.id), - ) - trial_policy = MagicMock(side_effect=AssertionError("get_app must not inspect trial policy")) - monkeypatch.setattr(RecommendedAppService, "is_trial_app_enabled", trial_policy) - - result = RecommendedAppService.get_app(app.id, session=sqlite_session) - - assert result is app - retrieval_instance.get_recommend_app_detail.assert_called_once_with(app.id, session=sqlite_session) - trial_policy.assert_not_called() - - def test_returns_none_when_app_is_not_recommended( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ) -> None: - app = _persist_app(sqlite_session, name="Private App") - - retrieval_instance = _mock_factory_for_app_detail(monkeypatch, result=None) - - result = RecommendedAppService.get_app(app.id, session=sqlite_session) - - assert result is None - retrieval_instance.get_recommend_app_detail.assert_called_once_with(app.id, session=sqlite_session) - - -# ── Pure logic tests: get_recommend_app_detail ───────────────────────── - - -class TestRecommendedAppServiceGetDetail: - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_returns_retrieval_detail_when_trial_disabled( - self, - mock_factory_class: MagicMock, - sqlite_session: Session, - ) -> None: - cases: list[tuple[str, RecommendedAppPayload]] = [ - ( - "complex-app", - _app_detail( - app_id="complex-app", - name="Complex App", - model_config={ - "provider": "openai", - "model": "gpt-4", - "parameters": {"temperature": 0.7, "max_tokens": 2000, "top_p": 1.0}, - }, - workflows=["workflow-1", "workflow-2"], - tools=["tool-1", "tool-2", "tool-3"], - ), - ), - ("app-empty", RecommendedAppPayload()), - ] - - for app_id, expected in cases: - mock_instance = MagicMock() - mock_instance.get_recommend_app_detail.return_value = expected - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - - result = RecommendedAppService.get_recommend_app_detail(app_id, session=sqlite_session) - - assert result is not None - assert result["can_trial"] is False - mock_instance.get_recommend_app_detail.assert_called_once_with(app_id, session=sqlite_session) - - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_different_modes( - self, - mock_factory_class: MagicMock, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - for mode in ["remote", "builtin", "db"]: - config_overrides(HOSTED_FETCH_APP_TEMPLATES_MODE=mode) - detail = _app_detail(app_id="test-app", name=f"App from {mode}") - mock_instance = MagicMock() - mock_instance.get_recommend_app_detail.return_value = detail - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - - result = RecommendedAppService.get_recommend_app_detail("test-app", session=sqlite_session) - - assert result is not None - mock_instance.get_recommend_app_detail.assert_called_with("test-app", session=sqlite_session) - mock_factory_class.get_recommend_app_factory.assert_called_with(mode) - - -# ── Pure logic tests: get_learn_dify_apps ────────────────────────────── - - -class TestRecommendedAppServiceGetLearnDifyApps: - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_uses_configured_retrieval_source( - self, - mock_factory_class: MagicMock, - sqlite_session: Session, - ) -> None: - expected_app = RecommendedAppPayload(app_id="app-1", category="Workflow") - mock_instance = MagicMock() - mock_instance.get_learn_dify_apps.return_value = { - "recommended_apps": [expected_app], - "categories": ["Workflow"], - } - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - - result = RecommendedAppService.get_learn_dify_apps("en-US", session=sqlite_session) - - assert result == {"recommended_apps": [{**expected_app, "can_trial": False}]} - mock_factory_class.get_recommend_app_factory.assert_called_once_with("remote") - mock_instance.get_learn_dify_apps.assert_called_once_with("en-US", session=sqlite_session) - - def test_sets_can_trial_when_trial_feature_enabled( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(HOSTED_FETCH_APP_TEMPLATES_MODE="db") - app = RecommendedAppPayload(app_id="app-1", category="Workflow") - mock_retrieval_instance = MagicMock() - mock_retrieval_instance.get_learn_dify_apps.return_value = { - "recommended_apps": [app], - "categories": ["Workflow"], - } - mock_retrieval_factory = MagicMock(return_value=mock_retrieval_instance) - monkeypatch.setattr( - service_module.RecommendAppRetrievalFactory, - "get_recommend_app_factory", - MagicMock(return_value=mock_retrieval_factory), - ) - monkeypatch.setattr(RecommendedAppService, "is_trial_app_enabled", MagicMock(return_value=True)) - trial_app_ids = MagicMock(return_value={"app-1"}) - monkeypatch.setattr(RecommendedAppService, "_get_trial_app_ids", trial_app_ids) - - result = RecommendedAppService.get_learn_dify_apps("en-US", session=sqlite_session) - - assert result["recommended_apps"][0]["can_trial"] is True - trial_app_ids.assert_called_once_with(sqlite_session, ["app-1"]) - - -# ── Integration tests: trial app features (real DB) ──────────────────── - - -class TestRecommendedAppServiceTrialFeatures: - def test_get_apps_should_not_query_trial_table_when_disabled( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ) -> None: - upstream_result = AppsResponse( - recommended_apps=[RecommendedAppPayload(app_id="app-1", can_trial=True)], categories=["all"] - ) - retrieval_instance, builtin_instance = _mock_factory_for_apps( - monkeypatch, mode="remote", result=upstream_result - ) - monkeypatch.setattr(RecommendedAppService, "is_trial_app_enabled", MagicMock(return_value=False)) - trial_app_ids = MagicMock(side_effect=AssertionError("disabled trial must not query TrialApp")) - monkeypatch.setattr(RecommendedAppService, "_get_trial_app_ids", trial_app_ids) - - result = RecommendedAppService.get_recommended_apps_and_categories("en-US", session=sqlite_session) - - assert result["recommended_apps"][0]["can_trial"] is False - trial_app_ids.assert_not_called() - retrieval_instance.get_recommended_apps_and_categories.assert_called_once_with("en-US", session=sqlite_session) - builtin_instance.fetch_recommended_apps_from_builtin.assert_not_called() - - def test_get_apps_should_enrich_can_trial_when_enabled( - self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch - ) -> None: - app_id_1 = str(uuid.uuid4()) - app_id_2 = str(uuid.uuid4()) - tenant_id = str(uuid.uuid4()) - - # app_id_1 has a TrialApp record; app_id_2 does not - sqlite_session.add(TrialApp(app_id=app_id_1, tenant_id=tenant_id)) - sqlite_session.commit() - - remote_result = AppsResponse(recommended_apps=[], categories=[]) - fallback_result = AppsResponse( - recommended_apps=[RecommendedAppPayload(app_id=app_id_1), RecommendedAppPayload(app_id=app_id_2)], - categories=["all"], - ) - _, builtin_instance = _mock_factory_for_apps( - monkeypatch, mode="remote", result=remote_result, fallback_result=fallback_result - ) - monkeypatch.setattr(RecommendedAppService, "is_trial_app_enabled", MagicMock(return_value=True)) - - result = RecommendedAppService.get_recommended_apps_and_categories("ja-JP", session=sqlite_session) - - builtin_instance.fetch_recommended_apps_from_builtin.assert_called_once_with("en-US") - assert result["recommended_apps"][0]["can_trial"] is True - assert result["recommended_apps"][1]["can_trial"] is False - - @pytest.mark.parametrize("has_trial_app", [True, False]) - def test_get_detail_should_set_can_trial_when_enabled( - self, - sqlite_session: Session, - monkeypatch: pytest.MonkeyPatch, - has_trial_app: bool, - ) -> None: - app_id = str(uuid.uuid4()) - tenant_id = str(uuid.uuid4()) - - if has_trial_app: - sqlite_session.add(TrialApp(app_id=app_id, tenant_id=tenant_id)) - sqlite_session.commit() - - detail = RecommendedAppPayload(id=app_id, name="Test App") - retrieval_instance = MagicMock() - retrieval_instance.get_recommend_app_detail.return_value = detail - retrieval_factory = MagicMock(return_value=retrieval_instance) - monkeypatch.setattr(service_module.dify_config, "HOSTED_FETCH_APP_TEMPLATES_MODE", "remote", raising=False) - monkeypatch.setattr( - service_module.RecommendAppRetrievalFactory, - "get_recommend_app_factory", - MagicMock(return_value=retrieval_factory), - ) - monkeypatch.setattr(RecommendedAppService, "is_trial_app_enabled", MagicMock(return_value=True)) - - result = RecommendedAppService.get_recommend_app_detail(app_id, session=sqlite_session) - assert result is not None - detail_result = cast(RecommendedAppPayload, result) - - assert detail_result["id"] == app_id - assert detail_result["can_trial"] is has_trial_app - - @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) - def test_get_detail_returns_none_before_reading_trial_flag( - self, - mock_factory_class: MagicMock, - sqlite_session: Session, - ) -> None: - mock_instance = MagicMock() - mock_instance.get_recommend_app_detail.return_value = None - mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - trial_policy = MagicMock(side_effect=AssertionError("missing app must not inspect trial policy")) - with patch.object(RecommendedAppService, "is_trial_app_enabled", trial_policy): - result = RecommendedAppService.get_recommend_app_detail("nonexistent", session=sqlite_session) - - assert result is None - mock_instance.get_recommend_app_detail.assert_called_once_with("nonexistent", session=sqlite_session) - trial_policy.assert_not_called() - - def test_add_trial_app_record_increments_count_for_existing(self, sqlite_session: Session) -> None: - app_id = str(uuid.uuid4()) - account_id = str(uuid.uuid4()) - - sqlite_session.add(AccountTrialAppRecord(app_id=app_id, account_id=account_id, count=3)) - sqlite_session.commit() - - RecommendedAppService.add_trial_app_record(app_id, account_id, session=sqlite_session) - - sqlite_session.expire_all() - record = sqlite_session.scalar( - select(AccountTrialAppRecord) - .where(AccountTrialAppRecord.app_id == app_id, AccountTrialAppRecord.account_id == account_id) - .limit(1) - ) - assert record is not None - assert record.count == 4 - - def test_add_trial_app_record_creates_new_record(self, sqlite_session: Session) -> None: - app_id = str(uuid.uuid4()) - account_id = str(uuid.uuid4()) - - RecommendedAppService.add_trial_app_record(app_id, account_id, session=sqlite_session) - - sqlite_session.expire_all() - record = sqlite_session.scalar( - select(AccountTrialAppRecord) - .where(AccountTrialAppRecord.app_id == app_id, AccountTrialAppRecord.account_id == account_id) - .limit(1) - ) - assert record is not None - assert record.app_id == app_id - assert record.account_id == account_id - assert record.count == 1 diff --git a/api/tests/unit_tests/services/test_snippet_dsl_service.py b/api/tests/unit_tests/services/test_snippet_dsl_service.py index a9bbdc98495..f213dd95267 100644 --- a/api/tests/unit_tests/services/test_snippet_dsl_service.py +++ b/api/tests/unit_tests/services/test_snippet_dsl_service.py @@ -533,12 +533,17 @@ def test_create_or_update_snippet_updates_existing_snippet_and_syncs_workflow(mo monkeypatch.setattr("services.snippet_dsl_service.SnippetService", lambda *_args, **_kwargs: snippet_service) monkeypatch.setattr( "services.snippet_dsl_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", - Mock(return_value=set()), + Mock(return_value={"retired-agent"}), ) monkeypatch.setattr( "services.snippet_dsl_service.WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync", Mock(), ) + retire_unowned = Mock() + monkeypatch.setattr( + "services.snippet_dsl_service.WorkflowAgentRetirementService.retire_unowned", + retire_unowned, + ) result = service._create_or_update_snippet( snippet=snippet, @@ -561,6 +566,11 @@ def test_create_or_update_snippet_updates_existing_snippet_and_syncs_workflow(mo assert snippet.icon_info == {"icon": "x"} snippet_service.sync_draft_workflow.assert_called_once() session.commit.assert_called_once() + retire_unowned.assert_called_once_with( + tenant_id="tenant-1", + agent_ids={"retired-agent"}, + account_id="account-1", + ) def test_create_or_update_snippet_creates_new_snippet_and_flushes(monkeypatch: pytest.MonkeyPatch): @@ -640,6 +650,40 @@ def test_export_snippet_dsl_returns_yaml(monkeypatch: pytest.MonkeyPatch): assert "input_fields:" in result +def test_export_snippet_dsl_uses_requested_published_workflow(monkeypatch: pytest.MonkeyPatch): + service = SnippetDslService(session=SimpleNamespace(get_bind=Mock())) + workflow = SimpleNamespace( + to_dict=Mock(return_value={"graph": {"nodes": []}}), + graph_dict={"nodes": []}, + ) + snippet = SimpleNamespace( + tenant_id="tenant-1", + name="Exported", + description=None, + type="node", + icon_info=None, + input_fields_list=[], + ) + get_published_workflow_by_id = Mock(return_value=workflow) + get_draft_workflow = Mock() + monkeypatch.setattr( + "services.snippet_dsl_service.SnippetService", + lambda *_args, **_kwargs: SimpleNamespace( + get_draft_workflow=get_draft_workflow, + get_published_workflow_by_id=get_published_workflow_by_id, + ), + ) + monkeypatch.setattr( + "services.snippet_dsl_service.DependenciesAnalysisService.generate_dependencies", + Mock(return_value=[]), + ) + + service.export_snippet_dsl(snippet, workflow_id="workflow-1") + + get_published_workflow_by_id.assert_called_once_with(snippet=snippet, workflow_id="workflow-1") + get_draft_workflow.assert_not_called() + + def test_append_workflow_export_data_filters_credentials_and_extracts_dependencies(monkeypatch: pytest.MonkeyPatch): service = SnippetDslService(session=SimpleNamespace()) workflow_dict = { diff --git a/api/tests/unit_tests/services/test_snippet_generate_service.py b/api/tests/unit_tests/services/test_snippet_generate_service.py index 83568a382ae..d1b3f428c7b 100644 --- a/api/tests/unit_tests/services/test_snippet_generate_service.py +++ b/api/tests/unit_tests/services/test_snippet_generate_service.py @@ -1,5 +1,4 @@ import json -from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import Mock @@ -7,7 +6,9 @@ import pytest from sqlalchemy.orm import Session, sessionmaker from core.workflow.snippet_start import SNIPPET_VIRTUAL_START_NODE_ID -from models.workflow import Workflow, WorkflowKind, WorkflowType +from models.account import Account +from models.snippet import CustomizedSnippet, SnippetType +from models.workflow import Workflow, WorkflowKind, WorkflowNodeExecutionModel, WorkflowType from services.snippet_generate_service import SnippetGenerateService @@ -28,8 +29,22 @@ def _workflow(graph: dict) -> Workflow: ) -def _session_maker(session: object | None = None) -> Mock: - return Mock(return_value=nullcontext(session or Mock())) +def _snippet(*, input_fields: list[dict] | None = None) -> CustomizedSnippet: + return CustomizedSnippet( + id="snippet-1", + tenant_id="tenant-1", + name="Snippet", + description="", + type=SnippetType.NODE, + created_by="account-1", + input_fields=json.dumps(input_fields) if input_fields else None, + ) + + +def _account(account_id: str = "user-1") -> Account: + account = Account(name="Test User", email=f"{account_id}@example.com") + account.id = account_id + return account def test_filter_virtual_start_events_keeps_blocking_response_unchanged(): @@ -67,7 +82,7 @@ def test_is_virtual_start_event(message, expected): def test_ensure_start_node_returns_workflow_when_start_already_exists(): workflow = _workflow({"nodes": [{"id": "start", "data": {"type": "start"}}], "edges": []}) - snippet = SimpleNamespace(input_fields_list=[]) + snippet = _snippet() result = SnippetGenerateService._ensure_start_node(workflow, snippet) @@ -83,8 +98,8 @@ def test_ensure_start_node_injects_virtual_start_for_root_candidates(monkeypatch "edges": [{"source": "llm-1", "target": "answer-1"}], } workflow = _workflow(graph) - snippet = SimpleNamespace( - input_fields_list=[ + snippet = _snippet( + input_fields=[ { "variable": "query", "label": "Query", @@ -139,8 +154,8 @@ def test_generate_raises_when_draft_workflow_missing(monkeypatch: pytest.MonkeyP with pytest.raises(ValueError, match="Workflow not initialized"): SnippetGenerateService.generate( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), - user=SimpleNamespace(id="user-1"), + snippet=_snippet(), + user=_account(), args={"inputs": {}}, invoke_from="debugger", ) @@ -148,8 +163,8 @@ def test_generate_raises_when_draft_workflow_missing(monkeypatch: pytest.MonkeyP def test_generate_delegates_to_workflow_generator_and_filters_stream(monkeypatch: pytest.MonkeyPatch): workflow = _workflow({"nodes": [{"id": "llm-1", "data": {"type": "llm"}}], "edges": []}) - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1", input_fields_list=[]) - user = SimpleNamespace(id="user-1") + snippet = _snippet() + user = _account() raw_stream = iter( [ {"event": "node_started", "data": {"node_id": SNIPPET_VIRTUAL_START_NODE_ID}}, @@ -189,8 +204,8 @@ def test_generate_delegates_to_workflow_generator_and_filters_stream(monkeypatch def test_run_published_delegates_to_workflow_generator_non_streaming(monkeypatch: pytest.MonkeyPatch): workflow = _workflow({"nodes": [{"id": "llm-1", "data": {"type": "llm"}}], "edges": []}) - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1", input_fields_list=[]) - user = SimpleNamespace(id="user-1") + snippet = _snippet() + user = _account() generator = SimpleNamespace(generate=Mock(return_value={"data": {"outputs": {"answer": "ok"}}})) monkeypatch.setattr( @@ -219,7 +234,7 @@ def test_run_published_delegates_to_workflow_generator_non_streaming(monkeypatch def test_ensure_start_node_for_worker_delegates(monkeypatch: pytest.MonkeyPatch): workflow = _workflow({"nodes": [], "edges": []}) - snippet = SimpleNamespace(input_fields_list=[]) + snippet = _snippet() ensure_start_node = Mock(return_value=workflow) monkeypatch.setattr(SnippetGenerateService, "_ensure_start_node", ensure_start_node) @@ -231,9 +246,9 @@ def test_ensure_start_node_for_worker_delegates(monkeypatch: pytest.MonkeyPatch) def test_run_draft_node_delegates_to_workflow_service(monkeypatch: pytest.MonkeyPatch): workflow = _workflow({"nodes": [{"id": "llm-1", "data": {"type": "llm"}}], "edges": []}) - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1") - account = SimpleNamespace(id="account-1") - execution = SimpleNamespace(id="execution-1") + snippet = _snippet() + account = _account("account-1") + execution = WorkflowNodeExecutionModel(id="execution-1") workflow_service = SimpleNamespace(run_draft_workflow_node=Mock(return_value=execution)) monkeypatch.setattr( @@ -271,10 +286,10 @@ def test_run_draft_node_raises_when_draft_workflow_missing(monkeypatch: pytest.M with pytest.raises(ValueError, match="Workflow not initialized"): SnippetGenerateService.run_draft_node( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), + snippet=_snippet(), node_id="llm-1", user_inputs={}, - account=SimpleNamespace(id="account-1"), + account=_account("account-1"), ) @@ -283,8 +298,8 @@ def test_generate_single_iteration_delegates_to_workflow_generator( sqlite_session_factory: sessionmaker[Session], ) -> None: workflow = _workflow({"nodes": [{"id": "iteration-1", "data": {"type": "iteration"}}], "edges": []}) - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1") - user = SimpleNamespace(id="user-1") + snippet = _snippet() + user = _account() response = iter(["event"]) generator = SimpleNamespace(single_iteration_generate=Mock(return_value=response)) workflow_generator_class = Mock(return_value=generator) @@ -316,7 +331,9 @@ def test_generate_single_iteration_delegates_to_workflow_generator( workflow_generator_class.convert_to_event_stream.assert_called_once_with(response) -def test_generate_single_iteration_raises_when_draft_workflow_missing(monkeypatch: pytest.MonkeyPatch): +def test_generate_single_iteration_raises_when_draft_workflow_missing( + monkeypatch: pytest.MonkeyPatch, unbound_session_factory: sessionmaker[Session] +): monkeypatch.setattr( "services.snippet_generate_service.SnippetService", lambda *_args, **_kwargs: SimpleNamespace(get_draft_workflow=Mock(return_value=None)), @@ -324,11 +341,11 @@ def test_generate_single_iteration_raises_when_draft_workflow_missing(monkeypatc with pytest.raises(ValueError, match="Workflow not initialized"): SnippetGenerateService.generate_single_iteration( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), - user=SimpleNamespace(id="user-1"), + snippet=_snippet(), + user=_account(), node_id="iteration-1", args={"inputs": {}}, - session_maker=_session_maker(), + session_maker=unbound_session_factory, ) @@ -337,8 +354,8 @@ def test_generate_single_loop_delegates_to_workflow_generator( sqlite_session_factory: sessionmaker[Session], ) -> None: workflow = _workflow({"nodes": [{"id": "loop-1", "data": {"type": "loop"}}], "edges": []}) - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1") - user = SimpleNamespace(id="user-1") + snippet = _snippet() + user = _account() response = iter(["event"]) generator = SimpleNamespace(single_loop_generate=Mock(return_value=response)) workflow_generator_class = Mock(return_value=generator) @@ -370,7 +387,9 @@ def test_generate_single_loop_delegates_to_workflow_generator( workflow_generator_class.convert_to_event_stream.assert_called_once_with(response) -def test_generate_single_loop_raises_when_draft_workflow_missing(monkeypatch: pytest.MonkeyPatch): +def test_generate_single_loop_raises_when_draft_workflow_missing( + monkeypatch: pytest.MonkeyPatch, unbound_session_factory: sessionmaker[Session] +): monkeypatch.setattr( "services.snippet_generate_service.SnippetService", lambda *_args, **_kwargs: SimpleNamespace(get_draft_workflow=Mock(return_value=None)), @@ -378,11 +397,11 @@ def test_generate_single_loop_raises_when_draft_workflow_missing(monkeypatch: py with pytest.raises(ValueError, match="Workflow not initialized"): SnippetGenerateService.generate_single_loop( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), - user=SimpleNamespace(id="user-1"), + snippet=_snippet(), + user=_account(), node_id="loop-1", args=SimpleNamespace(inputs={}), - session_maker=_session_maker(), + session_maker=unbound_session_factory, ) @@ -394,8 +413,8 @@ def test_run_published_raises_when_published_workflow_missing(monkeypatch: pytes with pytest.raises(ValueError, match="No published workflow found"): SnippetGenerateService.run_published( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), - user=SimpleNamespace(id="user-1"), + snippet=_snippet(), + user=_account(), args={"inputs": {}}, invoke_from="service-api", ) diff --git a/api/tests/unit_tests/services/test_snippet_service.py b/api/tests/unit_tests/services/test_snippet_service.py index 21a9deb824e..caac531bdbc 100644 --- a/api/tests/unit_tests/services/test_snippet_service.py +++ b/api/tests/unit_tests/services/test_snippet_service.py @@ -6,25 +6,36 @@ from types import SimpleNamespace from unittest.mock import Mock import pytest -from sqlalchemy import select +from sqlalchemy import event, select from sqlalchemy.orm import Session, sessionmaker from enums import DeploymentEdition from extensions.storage.storage_type import StorageType from graphon.variables.segments import StringSegment from graphon.variables.types import SegmentType -from models.agent import Agent, AgentScope, AgentSource, AgentStatus -from models.enums import CreatorUserRole -from models.model import UploadFile +from models.account import Account +from models.agent import ( + Agent, + AgentScope, + AgentSource, + AgentStatus, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) +from models.enums import AppStatus, CreatorUserRole +from models.model import App, AppMode, UploadFile from models.snippet import CustomizedSnippet, SnippetType from models.workflow import ( Workflow, WorkflowDraftVariable, WorkflowDraftVariableFile, WorkflowKind, + WorkflowNodeExecutionModel, + WorkflowRun, WorkflowType, ) from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError +from services.errors.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError from services.snippet_service import SnippetService @@ -45,21 +56,29 @@ def _create_workflow(*, workflow_id: str, version: str, graph: dict, features: d ) -def _snippet() -> CustomizedSnippet: - return CustomizedSnippet( - id="snippet-1", - tenant_id="tenant-1", - name="Snippet", - description="", - type=SnippetType.NODE, - created_by="account-1", - ) +def _snippet(**overrides) -> CustomizedSnippet: + values = { + "id": "snippet-1", + "tenant_id": "tenant-1", + "name": "Snippet", + "description": "", + "type": SnippetType.NODE, + "created_by": "account-1", + } + values.update(overrides) + return CustomizedSnippet(**values) + + +def _account(account_id: str = "account-1") -> Account: + account = Account(name="Test User", email=f"{account_id}@example.com") + account.id = account_id + return account def test_create_snippet_allows_duplicate_names( sqlite_session_factory: sessionmaker[Session], sqlite_session: Session ) -> None: - account = SimpleNamespace(id="account-1") + account = _account() existing = _snippet() existing.name = "shared name" sqlite_session.add(existing) @@ -207,8 +226,17 @@ def test_sync_draft_workflow_creates_draft_and_updates_input_fields( ) -> None: service = SnippetService(session_maker=sqlite_session_factory) monkeypatch.setattr(service, "get_draft_workflow", Mock(return_value=None)) + monkeypatch.setattr( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", + Mock(return_value={"retired-agent"}), + ) + retire_unowned = Mock() + monkeypatch.setattr( + "services.snippet_service.WorkflowAgentRetirementService.retire_unowned", + retire_unowned, + ) snippet = _snippet() - account = SimpleNamespace(id="account-1") + account = _account() workflow = service.sync_draft_workflow( snippet=snippet, @@ -227,20 +255,28 @@ def test_sync_draft_workflow_creates_draft_and_updates_input_fields( assert stored_workflow is not None assert stored_snippet is not None assert stored_snippet.input_fields_list == [{"variable": "query"}] + retire_unowned.assert_called_once_with( + tenant_id=snippet.tenant_id, + agent_ids={"retired-agent"}, + account_id=account.id, + ) def test_sync_draft_workflow_raises_when_hash_mismatches( sqlite_session_factory: sessionmaker[Session], ) -> None: service = SnippetService(session_maker=sqlite_session_factory) - service.get_draft_workflow = Mock(return_value=SimpleNamespace(unique_hash="server-hash")) + draft_workflow = _create_workflow( + workflow_id="workflow-1", version=Workflow.VERSION_DRAFT, graph={"nodes": []}, features={} + ) + service.get_draft_workflow = Mock(return_value=draft_workflow) with pytest.raises(WorkflowHashNotEqualError): service.sync_draft_workflow( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), + snippet=_snippet(), graph={"nodes": [], "edges": []}, unique_hash="client-hash", - account=SimpleNamespace(id="account-1"), + account=_account(), ) @@ -258,7 +294,7 @@ def test_sync_draft_workflow_updates_existing_draft_and_clears_variables( ) unique_hash = workflow.unique_hash snippet = _snippet() - account = SimpleNamespace(id="account-1") + account = _account() monkeypatch.setattr(service, "get_draft_workflow", Mock(return_value=workflow)) result = service.sync_draft_workflow( @@ -293,7 +329,7 @@ def test_update_workflow_updates_marked_fields(sqlite_session: Session) -> None: snippet = _snippet() sqlite_session.add_all([snippet, workflow]) sqlite_session.flush() - account = SimpleNamespace(id="account-1") + account = _account() result = service.update_workflow( session=sqlite_session, @@ -318,15 +354,73 @@ def test_update_workflow_returns_none_when_missing(sqlite_session: Session) -> N result = service.update_workflow( session=sqlite_session, - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), + snippet=_snippet(), workflow_id="missing-workflow", - account=SimpleNamespace(id="account-1"), + account=_account(), data={"marked_name": "v1"}, ) assert result is None +def test_delete_workflow_removes_published_version(sqlite_session: Session) -> None: + service = SnippetService.__new__(SnippetService) + workflow = _create_workflow( + workflow_id="workflow-1", + version="2026-01-01 00:00:00", + graph={"nodes": []}, + features={}, + ) + snippet = _snippet(workflow_id="workflow-2") + sqlite_session.add_all([snippet, workflow]) + sqlite_session.flush() + + result = service.delete_workflow(session=sqlite_session, snippet=snippet, workflow_id="workflow-1") + + assert result is True + sqlite_session.flush() + assert sqlite_session.get(Workflow, "workflow-1") is None + + +def test_delete_workflow_raises_when_missing(sqlite_session: Session) -> None: + service = SnippetService.__new__(SnippetService) + + with pytest.raises(ValueError, match="not found"): + service.delete_workflow(session=sqlite_session, snippet=_snippet(), workflow_id="missing-workflow") + + +def test_delete_workflow_raises_for_draft_version(sqlite_session: Session) -> None: + service = SnippetService.__new__(SnippetService) + workflow = _create_workflow( + workflow_id="workflow-1", + version=Workflow.VERSION_DRAFT, + graph={"nodes": []}, + features={}, + ) + snippet = _snippet() + sqlite_session.add_all([snippet, workflow]) + sqlite_session.flush() + + with pytest.raises(DraftWorkflowDeletionError): + service.delete_workflow(session=sqlite_session, snippet=snippet, workflow_id="workflow-1") + + +def test_delete_workflow_raises_when_currently_active(sqlite_session: Session) -> None: + service = SnippetService.__new__(SnippetService) + workflow = _create_workflow( + workflow_id="workflow-1", + version="2026-01-01 00:00:00", + graph={"nodes": []}, + features={}, + ) + snippet = _snippet(workflow_id="workflow-1") + sqlite_session.add_all([snippet, workflow]) + sqlite_session.flush() + + with pytest.raises(WorkflowInUseError): + service.delete_workflow(session=sqlite_session, snippet=snippet, workflow_id="workflow-1") + + def test_get_default_block_configs_skips_empty_defaults(monkeypatch: pytest.MonkeyPatch) -> None: node_with_default = SimpleNamespace(get_default_config=Mock(return_value={"type": "llm"})) node_without_default = SimpleNamespace(get_default_config=Mock(return_value=None)) @@ -375,7 +469,7 @@ def test_restore_published_snippet_workflow_to_draft_copies_source_snapshot( sqlite_session: Session, ) -> None: snippet = _snippet() - account = SimpleNamespace(id="account-2") + account = _account("account-2") source_graph = {"nodes": [{"id": "llm-1", "data": {"type": "llm"}}], "edges": []} source_features = {"opening_statement": "hello"} source_workflow = _create_workflow( @@ -394,6 +488,15 @@ def test_restore_published_snippet_workflow_to_draft_copies_source_snapshot( monkeypatch.setattr(service, "get_published_workflow_by_id", Mock(return_value=source_workflow)) monkeypatch.setattr(service, "get_draft_workflow", Mock(return_value=draft_workflow)) + monkeypatch.setattr( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.restore_agent_node_bindings_to_draft", + Mock(return_value={"retired-agent"}), + ) + retire_unowned = Mock() + monkeypatch.setattr( + "services.snippet_service.WorkflowAgentRetirementService.retire_unowned", + retire_unowned, + ) result = service.restore_published_workflow_to_draft( snippet=snippet, @@ -409,14 +512,19 @@ def test_restore_published_snippet_workflow_to_draft_copies_source_snapshot( stored = sqlite_session.get(Workflow, draft_workflow.id) assert stored is not None assert stored.graph_dict == source_graph + retire_unowned.assert_called_once_with( + tenant_id=snippet.tenant_id, + agent_ids={"retired-agent"}, + account_id=account.id, + ) def test_restore_published_snippet_workflow_to_draft_raises_when_source_missing( monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session], ) -> None: - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1") - account = SimpleNamespace(id="account-2") + snippet = _snippet() + account = _account("account-2") service = SnippetService(session_maker=sqlite_session_factory) monkeypatch.setattr(service, "get_published_workflow_by_id", Mock(return_value=None)) @@ -435,7 +543,7 @@ def test_restore_published_snippet_workflow_to_draft_adds_new_draft( sqlite_session: Session, ) -> None: snippet = _snippet() - account = SimpleNamespace(id="account-2") + account = _account("account-2") source_workflow = _create_workflow( workflow_id="published-workflow", version="2026-04-28 00:00:00", @@ -471,7 +579,7 @@ def test_restore_published_snippet_workflow_to_draft_adds_new_draft( def test_get_published_workflow_returns_none_without_workflow_id() -> None: service = SnippetService.__new__(SnippetService) - result = service.get_published_workflow(SimpleNamespace(id="snippet-1", tenant_id="tenant-1", workflow_id=None)) + result = service.get_published_workflow(_snippet()) assert result is None @@ -488,7 +596,7 @@ def test_get_published_workflow_by_id_raises_for_draft( with pytest.raises(IsDraftWorkflowError): service.get_published_workflow_by_id( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), + snippet=_snippet(), workflow_id="workflow-1", ) @@ -499,8 +607,8 @@ def test_publish_workflow_raises_when_draft_missing(sqlite_session: Session) -> with pytest.raises(ValueError, match="No valid workflow found"): service.publish_workflow( session=sqlite_session, - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), - account=SimpleNamespace(id="account-1"), + snippet=_snippet(), + account=_account(), ) @@ -534,10 +642,10 @@ def test_publish_workflow_creates_snapshot_and_updates_snippet( Mock(return_value=set()), ) - result, retirement_candidates = service.publish_workflow( + result = service.publish_workflow( session=sqlite_session, snippet=snippet, - account=SimpleNamespace(id="account-1"), + account=_account(), ) assert result.kind == WorkflowKind.SNIPPET @@ -548,7 +656,6 @@ def test_publish_workflow_creates_snapshot_and_updates_snippet( sqlite_session.flush() assert sqlite_session.get(Workflow, result.id) is result assert sqlite_session.get(CustomizedSnippet, snippet.id).workflow_id == result.id - assert retirement_candidates == set() def test_get_all_published_workflows_returns_empty_without_current_workflow(unbound_session: Session) -> None: @@ -556,7 +663,7 @@ def test_get_all_published_workflows_returns_empty_without_current_workflow(unbo result = service.get_all_published_workflows( session=unbound_session, - snippet=SimpleNamespace(id="snippet-1", workflow_id=None), + snippet=_snippet(), page=1, limit=20, ) @@ -580,7 +687,7 @@ def test_get_all_published_workflows_paginates(sqlite_session: Session) -> None: result, has_more = service.get_all_published_workflows( session=sqlite_session, - snippet=SimpleNamespace(id="snippet-1", workflow_id="workflow-current"), + snippet=_snippet(workflow_id="workflow-current"), page=1, limit=2, ) @@ -608,11 +715,12 @@ def test_delete_snippet_removes_related_records( assert observer.get(Workflow, workflow.id) is None -def test_delete_snippet_archives_owned_agents_and_schedules_backing_app_cleanup( +def test_delete_snippet_releases_last_owner_and_retries_archived_agent_cleanup_after_commit( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, ) -> None: snippet = _snippet() + archived_at = datetime(2025, 1, 1) agent = Agent( id="agent-1", tenant_id=snippet.tenant_id, @@ -623,13 +731,57 @@ def test_delete_snippet_archives_owned_agents_and_schedules_backing_app_cleanup( source=AgentSource.WORKFLOW, app_id=snippet.id, backing_app_id="backing-app-1", - status=AgentStatus.ACTIVE, - updated_by="creator-1", + workflow_id="workflow-1", + workflow_node_id="agent-node", + status=AgentStatus.ARCHIVED, + archived_by="original-account", + archived_at=archived_at, + updated_by="original-account", ) - sqlite_session.add_all([snippet, agent]) + hidden_app = App( + id="backing-app-1", + tenant_id=snippet.tenant_id, + name="Snippet Agent runtime", + mode=AppMode.AGENT, + status=AppStatus.NORMAL, + enable_site=False, + enable_api=False, + ) + workflow = _create_workflow( + workflow_id="workflow-1", + version=Workflow.VERSION_DRAFT, + graph={"nodes": []}, + features={}, + ) + owner_binding = WorkflowAgentNodeBinding( + id="snippet-inline-binding", + tenant_id=snippet.tenant_id, + app_id=snippet.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=agent.id, + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + sqlite_session.add_all([snippet, agent, hidden_app, workflow, owner_binding]) sqlite_session.flush() - cleanup_delay = Mock() - monkeypatch.setattr("tasks.remove_app_and_related_data_task.remove_app_and_related_data_task.delay", cleanup_delay) + hidden_app_id = hidden_app.id + workflow_id = workflow.id + owner_binding_id = owner_binding.id + events: list[str] = [] + event.listen(sqlite_session, "after_commit", lambda _session: events.append("commit"), once=True) + cleanup_delay = Mock(side_effect=lambda **_kwargs: events.append("cleanup-hidden-app")) + enqueue_collection = Mock(side_effect=lambda **_kwargs: events.append("enqueue-agent-purge")) + monkeypatch.setattr( + "services.agent.retirement_service.remove_app_and_related_data_task.delay", + cleanup_delay, + ) + monkeypatch.setattr( + "services.agent.retirement_service.enqueue_agent_resource_collection", + enqueue_collection, + ) result = SnippetService.delete_snippet( session=sqlite_session, @@ -638,13 +790,100 @@ def test_delete_snippet_archives_owned_agents_and_schedules_backing_app_cleanup( ) assert result is True - assert agent.status == "archived" - assert agent.archived_by == "account-1" - assert agent.archived_at is not None - assert agent.updated_by == "account-1" + assert agent.status == AgentStatus.ARCHIVED + assert ( + sqlite_session.scalar(select(WorkflowAgentNodeBinding).where(WorkflowAgentNodeBinding.id == owner_binding.id)) + is None + ) + cleanup_delay.assert_not_called() + enqueue_collection.assert_not_called() sqlite_session.commit() - assert sqlite_session.get(Agent, agent.id).status == AgentStatus.ARCHIVED + sqlite_session.expire_all() + stored_agent = sqlite_session.get(Agent, agent.id) + assert stored_agent is not None + assert stored_agent.status == AgentStatus.ARCHIVED + assert stored_agent.archived_by == "original-account" + assert stored_agent.archived_at == archived_at + assert stored_agent.updated_by == "original-account" + assert sqlite_session.get(App, hidden_app_id) is None + assert sqlite_session.get(Workflow, workflow_id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, owner_binding_id) is None + assert events == ["commit", "cleanup-hidden-app", "enqueue-agent-purge"] cleanup_delay.assert_called_once_with(tenant_id=snippet.tenant_id, app_id="backing-app-1") + enqueue_collection.assert_called_once_with( + tenant_id=snippet.tenant_id, + workspace_ids=[], + binding_ids=[], + home_snapshot_ids=[], + purge_agent_ids=[agent.id], + ) + + +def test_delete_snippet_keeps_agent_with_persisted_external_owner( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + snippet = _snippet() + agent = Agent( + id="agent-1", + tenant_id=snippet.tenant_id, + name="Shared workflow Agent", + description="", + role="", + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + app_id=snippet.id, + backing_app_id="backing-app-1", + status=AgentStatus.ACTIVE, + ) + workflow_app = App( + id="app-1", + tenant_id=snippet.tenant_id, + name="Workflow", + mode=AppMode.WORKFLOW, + status=AppStatus.NORMAL, + enable_site=True, + enable_api=True, + ) + workflow = Workflow( + id="workflow-1", + tenant_id=snippet.tenant_id, + app_id=workflow_app.id, + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph="{}", + features="{}", + created_by="account-1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + external_binding = WorkflowAgentNodeBinding( + tenant_id=snippet.tenant_id, + app_id=workflow_app.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=agent.id, + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + sqlite_session.add_all([snippet, agent, workflow_app, workflow, external_binding]) + sqlite_session.flush() + cleanup_delay = Mock() + monkeypatch.setattr( + "services.agent.retirement_service.remove_app_and_related_data_task.delay", + cleanup_delay, + ) + + SnippetService.delete_snippet(session=sqlite_session, snippet=snippet, account_id="account-1") + sqlite_session.commit() + sqlite_session.expire_all() + + assert sqlite_session.get(Agent, agent.id).status == AgentStatus.ACTIVE # type: ignore[union-attr] + assert sqlite_session.get(WorkflowAgentNodeBinding, external_binding.id) is not None + cleanup_delay.assert_not_called() def test_delete_draft_variable_files_removes_storage_objects( @@ -703,7 +942,7 @@ def test_delete_draft_variable_files_removes_storage_objects( def test_delete_archived_workflow_run_files_removes_prefixed_objects(monkeypatch: pytest.MonkeyPatch) -> None: from configs import dify_config - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1") + snippet = _snippet() archive_storage = SimpleNamespace( list_objects=Mock(return_value=["tenant-1/app_id=snippet-1/run.json"]), delete_object=Mock(), @@ -722,15 +961,15 @@ def test_workflow_run_queries_delegate_to_repositories(monkeypatch: pytest.Monke service = SnippetService.__new__(SnippetService) workflow_run_repo = SimpleNamespace( get_paginated_workflow_runs=Mock(return_value=SimpleNamespace(data=[])), - get_workflow_run_by_id=Mock(return_value=SimpleNamespace(id="run-1")), + get_workflow_run_by_id=Mock(return_value=WorkflowRun(id="run-1")), ) node_execution_repo = SimpleNamespace( - get_executions_by_workflow_run=Mock(return_value=[SimpleNamespace(id="node-execution-1")]), - get_node_last_execution=Mock(return_value=SimpleNamespace(id="last-run-1")), + get_executions_by_workflow_run=Mock(return_value=[WorkflowNodeExecutionModel(id="node-execution-1")]), + get_node_last_execution=Mock(return_value=WorkflowNodeExecutionModel(id="last-run-1")), ) service._workflow_run_repo = workflow_run_repo service._node_execution_service_repo = node_execution_repo - snippet = SimpleNamespace(id="snippet-1", tenant_id="tenant-1") + snippet = _snippet() expected_traces = [SimpleNamespace(id="node-execution-1:retry:1"), SimpleNamespace(id="node-execution-1")] mock_assemble = Mock(return_value=expected_traces) monkeypatch.setattr("services.snippet_service.assemble_workflow_node_execution_traces", mock_assemble) @@ -741,7 +980,9 @@ def test_workflow_run_queries_delegate_to_repositories(monkeypatch: pytest.Monke assert ( service.get_snippet_node_last_run( snippet=snippet, - workflow=SimpleNamespace(id="workflow-1"), + workflow=_create_workflow( + workflow_id="workflow-1", version=Workflow.VERSION_DRAFT, graph={"nodes": []}, features={} + ), node_id="llm-1", ).id == "last-run-1" @@ -774,7 +1015,7 @@ def test_workflow_run_node_executions_returns_empty_when_run_missing() -> None: service.get_snippet_workflow_run = Mock(return_value=None) result = service.get_snippet_workflow_run_node_executions( - snippet=SimpleNamespace(id="snippet-1", tenant_id="tenant-1"), + snippet=_snippet(), run_id="missing-run", ) diff --git a/api/tests/unit_tests/services/test_tag_application_service.py b/api/tests/unit_tests/services/test_tag_application_service.py new file mode 100644 index 00000000000..59921f98f4c --- /dev/null +++ b/api/tests/unit_tests/services/test_tag_application_service.py @@ -0,0 +1,49 @@ +from unittest.mock import MagicMock + +import pytest + +from machinery.context import RequestContext +from services.tag_application_service import ( + CreateTagInput, + TagApplicationService, + TagBindingInput, + TagSummary, + UpdateTagInput, +) + + +@pytest.fixture +def context() -> RequestContext: + return RequestContext("request-1", None, "account-1", "workspace-1") + + +def test_service_passes_stable_identity_to_store(context: RequestContext) -> None: + store = MagicMock() + store.list_tags.return_value = [TagSummary("tag-1", "Tag", "app", 1)] + store.create_tag.return_value = TagSummary("tag-2", "New", "app", 0) + store.update_tag.return_value = TagSummary("tag-2", "Updated", "app", 0) + service = TagApplicationService(tags=store) + + assert service.list_tags(context, "app", "search") == (TagSummary("tag-1", "Tag", "app", 1),) + service.create_tag(context, CreateTagInput("New", "app")) + service.update_tag(context, "tag-2", UpdateTagInput("Updated")) + service.delete_tag(context, "tag-2") + service.create_bindings(context, TagBindingInput(("tag-1",), "app-1", "app")) + service.delete_bindings(context, TagBindingInput(("tag-1",), "app-1", "app")) + + store.list_tags.assert_called_once_with("workspace-1", "app", "search") + store.create_tag.assert_called_once_with("workspace-1", "account-1", CreateTagInput("New", "app")) + store.update_tag.assert_called_once_with("workspace-1", "tag-2", UpdateTagInput("Updated")) + store.delete_tag.assert_called_once_with("workspace-1", "tag-2") + store.create_bindings.assert_called_once_with( + "workspace-1", "account-1", TagBindingInput(("tag-1",), "app-1", "app") + ) + store.delete_bindings.assert_called_once_with("workspace-1", TagBindingInput(("tag-1",), "app-1", "app")) + + +def test_service_rejects_context_without_active_workspace() -> None: + context = RequestContext("request-1", None, "account-1", None) + service = TagApplicationService(tags=MagicMock()) + + with pytest.raises(RuntimeError, match="active workspace"): + service.list_tags(context, "app") diff --git a/api/tests/unit_tests/services/test_vector_service.py b/api/tests/unit_tests/services/test_vector_service.py index 75d371c8fa9..50b0d2c58c4 100644 --- a/api/tests/unit_tests/services/test_vector_service.py +++ b/api/tests/unit_tests/services/test_vector_service.py @@ -16,7 +16,7 @@ import services.vector_service as vector_service_module from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from extensions.storage.storage_type import StorageType from models import UploadFile -from models.dataset import ChildChunk, DatasetProcessRule, SegmentAttachmentBinding +from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment, SegmentAttachmentBinding from models.dataset import Document as DatasetDocument from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, ProcessRuleMode from services.vector_service import VectorService @@ -42,16 +42,18 @@ def _make_dataset( is_multimodal: bool = False, embedding_model_provider: str | None = "openai", embedding_model: str = "text-embedding", -) -> MagicMock: - dataset = MagicMock(name="dataset") - dataset.id = dataset_id - dataset.tenant_id = tenant_id - dataset.doc_form = doc_form - dataset.indexing_technique = indexing_technique - dataset.is_multimodal = is_multimodal - dataset.embedding_model_provider = embedding_model_provider - dataset.embedding_model = embedding_model - dataset.get_doc_form.return_value = doc_form +) -> Dataset: + dataset = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Dataset", + created_by="account-1", + indexing_technique=indexing_technique, + chunk_structure=doc_form, + is_multimodal=is_multimodal, + embedding_model_provider=embedding_model_provider, + embedding_model=embedding_model, + ) return dataset @@ -64,24 +66,49 @@ def _make_segment( content: str = "hello", index_node_id: str = "node-1", index_node_hash: str = "hash-1", + session: Session | None = None, attachments: list[dict[str, str]] | None = None, -) -> MagicMock: - segment = MagicMock(name="segment") +) -> DocumentSegment: + segment = DocumentSegment( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + position=1, + content=content, + word_count=len(content), + tokens=len(content), + created_by="account-1", + index_node_id=index_node_id, + index_node_hash=index_node_hash, + ) segment.id = segment_id - segment.tenant_id = tenant_id - segment.dataset_id = dataset_id - segment.document_id = document_id - segment.content = content - segment.index_node_id = index_node_id - segment.index_node_hash = index_node_hash - segment.attachments = attachments or [] - segment.get_attachments.return_value = attachments or [] + if attachments: + assert session is not None + for attachment in attachments: + upload_file = _upload_file( + file_id=attachment["id"], + name=attachment.get("name", f"{attachment['id']}.png"), + tenant_id=tenant_id, + ) + session.add_all( + [ + upload_file, + SegmentAttachmentBinding( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + segment_id=segment_id, + attachment_id=upload_file.id, + ), + ] + ) + session.flush() return segment -def _upload_file(*, file_id: str = "file-1", name: str = "img.png") -> UploadFile: +def _upload_file(*, file_id: str = "file-1", name: str = "img.png", tenant_id: str = "tenant-1") -> UploadFile: upload_file = UploadFile( - tenant_id="tenant-1", + tenant_id=tenant_id, storage_type=StorageType.LOCAL, key=f"uploads/{file_id}", name=name, @@ -97,6 +124,30 @@ def _upload_file(*, file_id: str = "file-1", name: str = "img.png") -> UploadFil return upload_file +def _make_child_chunk( + *, + index_node_id: str, + content: str = "child", + index_node_hash: str = "hash", + tenant_id: str = "tenant-1", + dataset_id: str = "dataset-1", + document_id: str = "doc-1", + segment_id: str = "seg-1", +) -> ChildChunk: + return ChildChunk( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + segment_id=segment_id, + position=1, + content=content, + word_count=len(content), + created_by="account-1", + index_node_id=index_node_id, + index_node_hash=index_node_hash, + ) + + def test_create_segments_vector_regular_indexing_loads_documents_and_keywords( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: @@ -126,10 +177,11 @@ def test_create_segments_vector_regular_indexing_loads_multimodal_documents( ) -> None: dataset = _make_dataset(is_multimodal=True) segment = _make_segment( + session=sqlite_session, attachments=[ {"id": "img-1", "name": "a.png"}, {"id": "img-2", "name": "b.png"}, - ] + ], ) index_processor = MagicMock(name="index_processor") @@ -152,7 +204,7 @@ def test_create_segments_vector_regular_indexing_loads_multimodal_documents( assert second_args[1] == [] assert len(second_args[2]) == 2 assert second_kwargs["with_keywords"] is False - segment.get_attachments.assert_called_once_with(session=sqlite_session) + assert {document.page_content for document in second_args[2]} == {"a.png", "b.png"} def test_create_segments_vector_with_no_segments_does_not_load( @@ -171,7 +223,7 @@ def test_create_segments_vector_with_no_segments_does_not_load( def _persist_parent_child_rows( session: Session, *, - segment: MagicMock, + segment: DocumentSegment, include_document: bool = True, include_rule: bool = True, ) -> tuple[DatasetDocument | None, DatasetProcessRule | None]: @@ -415,13 +467,24 @@ def test_generate_child_chunks_regenerate_cleans_then_saves_children( dataset = _make_dataset(doc_form=IndexStructureType.PARAGRAPH_INDEX, tenant_id="tenant-1", dataset_id="dataset-1") segment = _make_segment(segment_id="seg-1") - dataset_document = MagicMock() - dataset_document.id = segment.document_id - dataset_document.doc_language = "en" - dataset_document.created_by = "user-1" - - processing_rule = MagicMock() - processing_rule.to_dict.return_value = {"rules": {}} + dataset_document = DatasetDocument( + id=segment.document_id, + tenant_id=segment.tenant_id, + dataset_id=segment.dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by="user-1", + doc_language="en", + ) + processing_rule = DatasetProcessRule( + dataset_id=segment.dataset_id, + mode=ProcessRuleMode.HIERARCHICAL, + rules="{}", + created_by="user-1", + ) child1 = _ChildDocStub(page_content="c1", metadata={"doc_id": "c1-id", "doc_hash": "c1-h"}) child2 = _ChildDocStub(page_content="c2", metadata={"doc_id": "c2-id", "doc_hash": "c2-h"}) @@ -456,12 +519,24 @@ def test_generate_child_chunks_flushes_even_when_no_children( ) -> None: dataset = _make_dataset(doc_form=IndexStructureType.PARAGRAPH_INDEX) segment = _make_segment() - dataset_document = MagicMock() - dataset_document.doc_language = "en" - dataset_document.created_by = "user-1" - - processing_rule = MagicMock() - processing_rule.to_dict.return_value = {"rules": {}} + dataset_document = DatasetDocument( + id=segment.document_id, + tenant_id=segment.tenant_id, + dataset_id=segment.dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by="user-1", + doc_language="en", + ) + processing_rule = DatasetProcessRule( + dataset_id=segment.dataset_id, + mode=ProcessRuleMode.HIERARCHICAL, + rules="{}", + created_by="user-1", + ) index_processor = MagicMock() index_processor.transform.return_value = [_ParentDocStub(children=[])] @@ -487,12 +562,7 @@ def test_create_child_chunk_vector_high_quality_adds_texts( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - child_chunk = MagicMock() - child_chunk.content = "child" - child_chunk.index_node_id = "id" - child_chunk.index_node_hash = "h" - child_chunk.document_id = "doc-1" - child_chunk.dataset_id = "dataset-1" + child_chunk = _make_child_chunk(index_node_id="id", index_node_hash="h") vector_instance = MagicMock() vector_cls = MagicMock(return_value=vector_instance) @@ -508,12 +578,7 @@ def test_create_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch, vector_cls = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - child_chunk = MagicMock() - child_chunk.content = "child" - child_chunk.index_node_id = "id" - child_chunk.index_node_hash = "h" - child_chunk.document_id = "doc-1" - child_chunk.dataset_id = "dataset-1" + child_chunk = _make_child_chunk(index_node_id="id", index_node_hash="h") VectorService.create_child_chunk_vector(child_chunk, dataset, session=sqlite_session) vector_cls.assert_not_called() @@ -524,22 +589,13 @@ def test_update_child_chunk_vector_high_quality_updates_vector( ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - new_chunk = MagicMock() - new_chunk.content = "n" - new_chunk.index_node_id = "nid" - new_chunk.index_node_hash = "nh" - new_chunk.document_id = "d" - new_chunk.dataset_id = "ds" - - upd_chunk = MagicMock() - upd_chunk.content = "u" - upd_chunk.index_node_id = "uid" - upd_chunk.index_node_hash = "uh" - upd_chunk.document_id = "d" - upd_chunk.dataset_id = "ds" - - del_chunk = MagicMock() - del_chunk.index_node_id = "did" + new_chunk = _make_child_chunk( + content="n", index_node_id="nid", index_node_hash="nh", document_id="d", dataset_id="ds" + ) + upd_chunk = _make_child_chunk( + content="u", index_node_id="uid", index_node_hash="uh", document_id="d", dataset_id="ds" + ) + del_chunk = _make_child_chunk(index_node_id="did") vector_instance = MagicMock() vector_cls = MagicMock(return_value=vector_instance) @@ -564,8 +620,7 @@ def test_update_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch, def test_delete_child_chunk_vector_deletes_by_id(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: dataset = _make_dataset() - child_chunk = MagicMock() - child_chunk.index_node_id = "cid" + child_chunk = _make_child_chunk(index_node_id="cid") vector_instance = MagicMock() vector_cls = MagicMock(return_value=vector_instance) @@ -585,7 +640,7 @@ def test_update_multimodel_vector_returns_when_not_high_quality( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY, is_multimodal=True) - segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}]) + segment = _make_segment(tenant_id="t") vector_cls = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", vector_cls) @@ -601,7 +656,7 @@ def test_update_multimodel_vector_returns_when_no_actual_change( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}, {"id": "b"}]) + segment = _make_segment(tenant_id="t", session=sqlite_session, attachments=[{"id": "a"}, {"id": "b"}]) vector_cls = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", vector_cls) @@ -610,7 +665,7 @@ def test_update_multimodel_vector_returns_when_no_actual_change( session=sqlite_session, segment=segment, attachment_ids=["b", "a"], dataset=dataset ) vector_cls.assert_not_called() - assert not sqlite_session.in_transaction() + assert sqlite_session.in_transaction() def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids( @@ -618,23 +673,14 @@ def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids( sqlite_session: Session, ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}, {"id": "old-2"}]) + segment = _make_segment( + tenant_id="tenant-1", + session=sqlite_session, + attachments=[{"id": "old-1"}, {"id": "old-2"}], + ) vector_instance = MagicMock(name="vector_instance") vector_cls = MagicMock(return_value=vector_instance) - sqlite_session.add_all( - [ - SegmentAttachmentBinding( - tenant_id="tenant-1", - dataset_id="dataset-1", - document_id="doc-1", - segment_id="seg-1", - attachment_id=attachment_id, - ) - for attachment_id in ("old-1", "old-2") - ] - ) - sqlite_session.flush() monkeypatch.setattr(vector_service_module, "Vector", vector_cls) @@ -650,7 +696,7 @@ def test_update_multimodel_vector_flushes_when_no_upload_files_found( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}]) + segment = _make_segment(tenant_id="tenant-1", session=sqlite_session, attachments=[{"id": "old-1"}]) vector_instance = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) @@ -668,7 +714,12 @@ def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_up sqlite_session: Session, ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}]) + segment = _make_segment( + segment_id="seg-1", + tenant_id="tenant-1", + session=sqlite_session, + attachments=[{"id": "old-1"}], + ) vector_instance = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) @@ -696,7 +747,7 @@ def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops sqlite_session: Session, ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=False) - segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}]) + segment = _make_segment(tenant_id="tenant-1", session=sqlite_session, attachments=[{"id": "old-1"}]) vector_instance = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) @@ -719,7 +770,12 @@ def test_update_multimodel_vector_rolls_back_and_reraises_on_error( sqlite_session: Session, ) -> None: dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}]) + segment = _make_segment( + segment_id="seg-1", + tenant_id="tenant-1", + session=sqlite_session, + attachments=[{"id": "old-1"}], + ) vector_instance = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) @@ -730,10 +786,11 @@ def test_update_multimodel_vector_rolls_back_and_reraises_on_error( monkeypatch.setattr(sqlite_session, "flush", MagicMock(side_effect=RuntimeError("boom"))) with caplog.at_level(logging.ERROR, logger="services.vector_service"): - with pytest.raises(RuntimeError, match="boom"): - VectorService.update_multimodel_vector( - session=sqlite_session, segment=segment, attachment_ids=["file-1"], dataset=dataset - ) + with sqlite_session.no_autoflush: + with pytest.raises(RuntimeError, match="boom"): + VectorService.update_multimodel_vector( + session=sqlite_session, segment=segment, attachment_ids=["file-1"], dataset=dataset + ) assert any(r.levelno >= logging.ERROR for r in caplog.records) assert rollback_events == ["rollback"] diff --git a/api/tests/unit_tests/services/test_web_app_runtime_query_service.py b/api/tests/unit_tests/services/test_web_app_runtime_query_service.py new file mode 100644 index 00000000000..5464c33ec91 --- /dev/null +++ b/api/tests/unit_tests/services/test_web_app_runtime_query_service.py @@ -0,0 +1,144 @@ +from unittest.mock import MagicMock, create_autospec + +import pytest + +from services.app_definition_query_service import AppSiteConfiguration +from services.entities.feature_entities import FeatureModel +from services.file_service import FileService +from services.web_app_runtime_query_service import ( + WebAppBootstrap, + WebAppRuntimeQuery, + WebAppRuntimeQueryService, + WebAppRuntimeRecord, + WebAppRuntimeUnavailableError, +) + +_FILES_URL = "https://files.example.com" + + +@pytest.fixture +def workspace_features() -> MagicMock: + return MagicMock(return_value=FeatureModel()) + + +def _site_configuration() -> AppSiteConfiguration: + return AppSiteConfiguration( + title="Test Site", + chat_color_theme="light", + chat_color_theme_inverted=False, + icon_type="image", + icon="file-1", + icon_background="#ffffff", + description="Description", + copyright="Copyright", + privacy_policy="Privacy", + input_placeholder="Ask anything", + custom_disclaimer="Disclaimer", + default_language="en-US", + prompt_public=True, + show_workflow_steps=True, + use_icon_as_answer_icon=False, + ) + + +def _runtime_record( + *, + mode: str = "agent-chat", + tenant_status: str = "normal", + tenant_custom_config_json: str | None = '{"remove_webapp_brand":true,"replace_webapp_logo":"file-2"}', +) -> WebAppRuntimeRecord: + return WebAppRuntimeRecord( + app_id="app-1", + tenant_id="tenant-1", + mode=mode, + enable_site=True, + site=_site_configuration(), + plan="pro", + tenant_status=tenant_status, + tenant_custom_config_json=tenant_custom_config_json, + ) + + +def _service( + runtime: MagicMock, + *, + file_service: MagicMock | None = None, + workspace_features: MagicMock | None = None, +) -> WebAppRuntimeQueryService: + if file_service is None: + file_service = MagicMock(spec=FileService) + file_service.get_icon_url.return_value = None + if workspace_features is None: + workspace_features = MagicMock(return_value=FeatureModel()) + return WebAppRuntimeQueryService( + runtime=runtime, + file_service=file_service, + workspace_features=workspace_features, + files_url=_FILES_URL, + ) + + +@pytest.mark.parametrize("record", [None, _runtime_record(tenant_status="archive")]) +def test_get_bootstrap_rejects_unavailable_runtime(record: WebAppRuntimeRecord | None) -> None: + runtime: MagicMock = create_autospec(WebAppRuntimeQuery, instance=True, spec_set=True) + runtime.get_runtime_record.return_value = record + + with pytest.raises(WebAppRuntimeUnavailableError, match="Site not found"): + _service(runtime).get_bootstrap("app-1") + + +def test_get_bootstrap_applies_feature_and_branding_policy_after_record_load( + workspace_features: MagicMock, +) -> None: + runtime: MagicMock = create_autospec(WebAppRuntimeQuery, instance=True, spec_set=True) + record = _runtime_record() + features = FeatureModel(can_replace_logo=True, webapp_copyright_enabled=False) + features.billing.enabled = True + events: list[str] = [] + runtime.get_runtime_record.side_effect = lambda _app_id: events.append("record") or record + workspace_features.side_effect = lambda _tenant_id, **_kwargs: events.append("features") or features + file_service = MagicMock(spec=FileService) + file_service.get_icon_url.side_effect = lambda *_args, **_kwargs: events.append("icon") or "https://icon" + + result = _service( + runtime, + file_service=file_service, + workspace_features=workspace_features, + ).get_bootstrap("app-1") + + assert result == WebAppBootstrap( + app_id="app-1", + mode="agent-chat", + enable_site=True, + site={ + **record.site._asdict(), + "copyright": None, + "input_placeholder": None, + "icon_url": "https://icon", + }, + plan="pro", + can_replace_logo=True, + custom_config={ + "remove_webapp_brand": True, + "replace_webapp_logo": "https://files.example.com/files/workspaces/tenant-1/webapp-logo", + }, + ) + assert events == ["record", "features", "icon"] + workspace_features.assert_called_once_with("tenant-1") + file_service.get_icon_url.assert_called_once_with("file-1", "tenant-1") + + +def test_get_bootstrap_skips_legacy_custom_config_when_branding_is_not_allowed( + workspace_features: MagicMock, +) -> None: + runtime: MagicMock = create_autospec(WebAppRuntimeQuery, instance=True, spec_set=True) + record = _runtime_record(tenant_custom_config_json="not-json") + runtime.get_runtime_record.return_value = record + + workspace_features.return_value = FeatureModel(can_replace_logo=False) + + result = _service(runtime, workspace_features=workspace_features).get_bootstrap("app-1") + + assert result.site == {**record.site._asdict(), "icon_url": None} + assert result.can_replace_logo is False + assert result.custom_config is None diff --git a/api/tests/unit_tests/services/test_webapp_access_query_service.py b/api/tests/unit_tests/services/test_webapp_access_query_service.py index 8c0702a9018..36007fccaeb 100644 --- a/api/tests/unit_tests/services/test_webapp_access_query_service.py +++ b/api/tests/unit_tests/services/test_webapp_access_query_service.py @@ -16,21 +16,25 @@ def _service( access: MagicMock, enabled: bool = True, access_mode: WebAppAccessMode = WebAppAccessMode.PRIVATE, -) -> tuple[WebAppAccessQueryService, MagicMock]: + allowed: bool = True, +) -> tuple[WebAppAccessQueryService, MagicMock, MagicMock]: access_mode_for_app = MagicMock(return_value=access_mode) + is_user_allowed_for_app = MagicMock(return_value=allowed) return ( WebAppAccessQueryService( access=access, webapp_auth_enabled=enabled, access_mode_for_app=access_mode_for_app, + is_user_allowed_for_app=is_user_allowed_for_app, ), access_mode_for_app, + is_user_allowed_for_app, ) def test_disabled_auth_returns_public_before_resolving_app() -> None: access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) - service, access_mode_for_app = _service(access=access, enabled=False) + service, access_mode_for_app, _ = _service(access=access, enabled=False) assert service.get_access_mode(app_id=None, app_code=None) is WebAppAccessMode.PUBLIC access.find_app_id_by_code.assert_not_called() @@ -39,7 +43,7 @@ def test_disabled_auth_returns_public_before_resolving_app() -> None: def test_enabled_auth_reads_access_mode_by_app_id() -> None: access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) - service, access_mode_for_app = _service(access=access) + service, access_mode_for_app, _ = _service(access=access) assert service.get_access_mode(app_id="app-1", app_code=None) is WebAppAccessMode.PRIVATE access.find_app_id_by_code.assert_not_called() @@ -49,7 +53,7 @@ def test_enabled_auth_reads_access_mode_by_app_id() -> None: def test_app_code_takes_precedence_over_app_id() -> None: access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) access.find_app_id_by_code.return_value = "resolved-id" - service, access_mode_for_app = _service(access=access, access_mode=WebAppAccessMode.SSO_VERIFIED) + service, access_mode_for_app, _ = _service(access=access, access_mode=WebAppAccessMode.SSO_VERIFIED) assert service.get_access_mode(app_id="ignored-id", app_code="code-1") is WebAppAccessMode.SSO_VERIFIED access.find_app_id_by_code.assert_called_once_with("code-1") @@ -59,7 +63,7 @@ def test_app_code_takes_precedence_over_app_id() -> None: def test_missing_app_code_raises_not_found() -> None: access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) access.find_app_id_by_code.return_value = None - service, access_mode_for_app = _service(access=access) + service, access_mode_for_app, _ = _service(access=access) with pytest.raises(WebAppAccessAppNotFoundError): service.get_access_mode(app_id="must-not-fallback", app_code="missing-code") @@ -69,7 +73,7 @@ def test_missing_app_code_raises_not_found() -> None: def test_enabled_auth_requires_app_id_or_code() -> None: access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) - service, access_mode_for_app = _service(access=access) + service, access_mode_for_app, _ = _service(access=access) with pytest.raises(WebAppAccessReferenceRequiredError, match="^appId or appCode must be provided$"): service.get_access_mode(app_id=None, app_code=None) @@ -81,7 +85,7 @@ def test_repository_failure_is_not_hidden() -> None: access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) failure = TypeError("repository bug") access.find_app_id_by_code.side_effect = failure - service, _ = _service(access=access) + service, _, _ = _service(access=access) with pytest.raises(TypeError) as raised: service.get_access_mode(app_id=None, app_code="code-1") @@ -91,7 +95,7 @@ def test_repository_failure_is_not_hidden() -> None: def test_access_mode_failure_is_not_hidden() -> None: access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) - service, access_mode_for_app = _service(access=access) + service, access_mode_for_app, _ = _service(access=access) failure = TypeError("adapter bug") access_mode_for_app.side_effect = failure @@ -99,3 +103,48 @@ def test_access_mode_failure_is_not_hidden() -> None: service.get_access_mode(app_id="app-1", app_code=None) assert raised.value is failure + + +@pytest.mark.parametrize( + ("access_mode", "expected"), + [ + pytest.param(WebAppAccessMode.PUBLIC, False, id="public"), + pytest.param(WebAppAccessMode.SSO_VERIFIED, False, id="sso-verified"), + pytest.param(WebAppAccessMode.PRIVATE, True, id="private"), + pytest.param(WebAppAccessMode.PRIVATE_ALL, True, id="private-all"), + ], +) +def test_requires_permission_check_for_private_modes(access_mode: WebAppAccessMode, expected: bool) -> None: + access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) + service, access_mode_for_app, _ = _service(access=access, access_mode=access_mode) + + assert service.requires_permission_check("app-1") is expected + access_mode_for_app.assert_called_once_with("app-1") + + +def test_disabled_auth_still_reads_configured_mode_before_passport() -> None: + access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) + service, access_mode_for_app, _ = _service( + access=access, + enabled=False, + access_mode=WebAppAccessMode.PRIVATE, + ) + + assert service.requires_permission_check("app-1") is True + access_mode_for_app.assert_called_once_with("app-1") + + +def test_disabled_auth_allows_after_passport_without_querying_user_permission() -> None: + access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) + service, _, is_user_allowed_for_app = _service(access=access, enabled=False, allowed=False) + + assert service.is_user_allowed(user_id="user-1", app_id="app-1") is True + is_user_allowed_for_app.assert_not_called() + + +def test_enabled_auth_delegates_user_permission() -> None: + access: MagicMock = create_autospec(WebAppAccessQuery, instance=True, spec_set=True) + service, _, is_user_allowed_for_app = _service(access=access, allowed=False) + + assert service.is_user_allowed(user_id="user-1", app_id="app-1") is False + is_user_allowed_for_app.assert_called_once_with("user-1", "app-1") diff --git a/api/tests/unit_tests/services/test_workflow_collaboration_service.py b/api/tests/unit_tests/services/test_workflow_collaboration_service.py index 77e55b545c3..3cf1f3c0a21 100644 --- a/api/tests/unit_tests/services/test_workflow_collaboration_service.py +++ b/api/tests/unit_tests/services/test_workflow_collaboration_service.py @@ -8,6 +8,7 @@ from socketio.exceptions import TimeoutError as SocketIOTimeoutError from sqlalchemy import Engine from sqlalchemy.orm import Session +from core.rbac import RBACPermission, RBACResourceScope from models.account import Account, Tenant from models.base import TypeBase from models.model import App, AppMode, IconType @@ -24,7 +25,7 @@ def db_session(sqlite_engine: Engine) -> Iterator[Session]: yield session -def _app(*, app_id: str, tenant_id: str) -> App: +def _app(*, app_id: str, tenant_id: str, maintainer: str | None = None) -> App: return App( id=app_id, tenant_id=tenant_id, @@ -42,6 +43,7 @@ def _app(*, app_id: str, tenant_id: str) -> App: is_public=False, is_universal=False, max_active_requests=None, + maintainer=maintainer, use_icon_as_answer_icon=False, ) @@ -65,23 +67,37 @@ class TestWorkflowCollaborationService: "avatar": None, "tenant_id": "t-1", } + db_session.add(_app(app_id="wf-1", tenant_id="t-1", maintainer="owner-1")) + db_session.commit() with ( - patch.object(collaboration_service, "_can_access_workflow", return_value=True), - patch.object(collaboration_service, "get_or_set_leader", return_value="sid-1"), - patch.object(collaboration_service, "broadcast_online_users"), + patch("services.workflow_collaboration_service.dify_config.RBAC_ENABLED", True), + patch( + "services.workflow_collaboration_service.RBACService.CheckAccess.check", return_value=True + ) as check_access, + patch.object(collaboration_service, "get_or_set_leader", return_value="sid-1") as get_leader, + patch.object(collaboration_service, "broadcast_online_users") as broadcast_online_users, ): # Act result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1", session=db_session) # Assert assert result == ("u-1", True) + check_access.assert_called_once_with( + "t-1", + "u-1", + scene=RBACPermission.APP_EDIT, + resource_type=RBACResourceScope.APP, + resource_id="wf-1", + ) repository.set_session_info.assert_called_once() session_info = repository.set_session_info.call_args.args[1] assert session_info["server_id"] == "server-1" repository.refresh_server_heartbeat.assert_called_once_with("server-1") socketio.start_background_task.assert_called_once() + get_leader.assert_called_once_with("wf-1", "sid-1") socketio.enter_room.assert_called_once_with("sid-1", "wf-1") + broadcast_online_users.assert_called_once_with("wf-1") socketio.emit.assert_called_once_with("status", {"isLeader": True}, room="sid-1") def test_authorize_and_join_workflow_room_returns_none_when_missing_user( @@ -120,13 +136,33 @@ class TestWorkflowCollaborationService: "avatar": None, "tenant_id": "t-1", } + db_session.add(_app(app_id="wf-1", tenant_id="t-1", maintainer="owner-1")) + db_session.commit() - with patch.object(collaboration_service, "_can_access_workflow", return_value=False): + with ( + patch("services.workflow_collaboration_service.dify_config.RBAC_ENABLED", True), + patch( + "services.workflow_collaboration_service.RBACService.CheckAccess.check", return_value=False + ) as check_access, + patch.object(collaboration_service, "get_or_set_leader") as get_leader, + patch.object(collaboration_service, "broadcast_online_users") as broadcast_online_users, + ): result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1", session=db_session) assert result is None + check_access.assert_called_once_with( + "t-1", + "u-1", + scene=RBACPermission.APP_EDIT, + resource_type=RBACResourceScope.APP, + resource_id="wf-1", + ) + repository.refresh_server_heartbeat.assert_not_called() repository.set_session_info.assert_not_called() + socketio.start_background_task.assert_not_called() + get_leader.assert_not_called() socketio.enter_room.assert_not_called() + broadcast_online_users.assert_not_called() socketio.emit.assert_not_called() def test_repr_and_save_socket_identity(self, service: tuple[WorkflowCollaborationService, Mock, Mock]) -> None: @@ -159,11 +195,34 @@ class TestWorkflowCollaborationService: ) db_session.commit() - result = collaboration_service._can_access_workflow("wf-1", "tenant-1", session=db_session) + with patch("services.workflow_collaboration_service.dify_config.RBAC_ENABLED", False): + result = collaboration_service._can_access_workflow("wf-1", "tenant-1", "user-1", session=db_session) + + assert result is True + assert ( + collaboration_service._can_access_workflow("wf-1", "tenant-other", "user-1", session=db_session) + is False + ) + assert ( + collaboration_service._can_access_workflow("wf-other", "tenant-other", "user-1", session=db_session) + is True + ) + + def test_can_access_workflow_allows_maintainer_without_rbac_call( + self, service: tuple[WorkflowCollaborationService, Mock, Mock], db_session: Session + ) -> None: + collaboration_service, _repository, _socketio = service + db_session.add(_app(app_id="wf-1", tenant_id="tenant-1", maintainer="owner-1")) + db_session.commit() + + with ( + patch("services.workflow_collaboration_service.dify_config.RBAC_ENABLED", True), + patch("services.workflow_collaboration_service.RBACService.CheckAccess.check") as check_access, + ): + result = collaboration_service._can_access_workflow("wf-1", "tenant-1", "owner-1", session=db_session) assert result is True - assert collaboration_service._can_access_workflow("wf-1", "tenant-other", session=db_session) is False - assert collaboration_service._can_access_workflow("wf-other", "tenant-other", session=db_session) is True + check_access.assert_not_called() def test_relay_collaboration_event_unauthorized( self, service: tuple[WorkflowCollaborationService, Mock, Mock] diff --git a/api/tests/unit_tests/services/test_workflow_service.py b/api/tests/unit_tests/services/test_workflow_service.py index 4c6e223719f..6b77c96e568 100644 --- a/api/tests/unit_tests/services/test_workflow_service.py +++ b/api/tests/unit_tests/services/test_workflow_service.py @@ -39,10 +39,22 @@ from graphon.variables import StringVariable from graphon.variables.input_entities import VariableEntityType from libs.datetime_utils import naive_utc_now from models.account import Account +from models.agent import ( + Agent, + AgentConfigSnapshot, + AgentKind, + AgentScope, + AgentSource, + AgentStatus, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) +from models.agent_config_entities import AgentSoulConfig from models.human_input import HumanInputFormRecipient, RecipientType from models.model import App, AppMode from models.tools import BuiltinToolProvider, WorkflowToolProvider from models.workflow import Workflow, WorkflowType +from services.agent.retirement_service import WorkflowAgentRetirementService from services.errors.app import IsDraftWorkflowError, TriggerNodeLimitExceededError, WorkflowHashNotEqualError from services.errors.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError from services.workflow_ref_service import WorkflowRef @@ -219,6 +231,26 @@ class TestWorkflowService: """Create a WorkflowService whose repositories use the test SQLite engine.""" return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + def test_get_tenant_app_maintainers_scopes_requested_apps( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + sqlite_session.add_all( + [ + TestWorkflowAssociatedDataFactory.create_app( + app_id="app-1", tenant_id="tenant-1", maintainer="owner-1" + ), + TestWorkflowAssociatedDataFactory.create_app(app_id="app-2", tenant_id="tenant-1"), + TestWorkflowAssociatedDataFactory.create_app( + app_id="app-3", tenant_id="tenant-2", maintainer="owner-2" + ), + ] + ) + sqlite_session.commit() + + assert workflow_service.get_tenant_app_maintainers( + ["app-1", "app-2", "app-3", "missing"], "tenant-1", session=sqlite_session + ) == {"app-1": "owner-1", "app-2": None} + # ==================== Workflow Existence Tests ==================== # These tests verify the service can check if a draft workflow exists @@ -326,6 +358,24 @@ class TestWorkflowService: assert result is workflow + def test_get_published_workflow_by_id_can_lock_restore_source(self, workflow_service: WorkflowService): + app = TestWorkflowAssociatedDataFactory.create_app() + workflow = TestWorkflowAssociatedDataFactory.create_workflow(version="v1") + session = MagicMock(spec=Session) + session.scalar.return_value = workflow + + result = workflow_service.get_published_workflow_by_id( + app, + workflow.id, + session=session, + for_update=True, + ) + + stmt = session.scalar.call_args.args[0] + sql = str(stmt.compile(dialect=postgresql.dialect())) + assert result is workflow + assert "FOR UPDATE" in sql + def test_get_published_workflow_by_id_raises_error_for_draft( self, workflow_service: WorkflowService, sqlite_session: Session ): @@ -419,7 +469,14 @@ class TestWorkflowService: graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() features = {"file_upload": {"enabled": False}} - with patch("services.workflow_service.app_draft_workflow_was_synced"): + with ( + patch("services.workflow_service.app_draft_workflow_was_synced"), + patch( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", + return_value={"retired-agent"}, + ), + patch("services.workflow_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + ): result = workflow_service.sync_draft_workflow( app_model=app, graph=graph, @@ -435,6 +492,11 @@ class TestWorkflowService: assert persisted_workflow is result assert result.graph_dict == graph assert result.features_dict == features + retire_unowned.assert_called_once_with( + tenant_id=app.tenant_id, + agent_ids={"retired-agent"}, + account_id=account.id, + ) def test_sync_draft_workflow_updates_existing_draft( self, workflow_service: WorkflowService, sqlite_session: Session @@ -759,7 +821,19 @@ class TestWorkflowService: sqlite_session.add_all([source_workflow, draft_workflow]) sqlite_session.commit() - with patch("services.workflow_service.app_draft_workflow_was_synced"): + with ( + patch("services.workflow_service.app_draft_workflow_was_synced"), + patch.object( + workflow_service, + "get_published_workflow_by_id", + wraps=workflow_service.get_published_workflow_by_id, + ) as get_published_workflow_by_id, + patch( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.restore_agent_node_bindings_to_draft", + return_value={"retired-agent"}, + ), + patch("services.workflow_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + ): result = workflow_service.restore_published_workflow_to_draft( app_model=app, workflow_id=source_workflow.id, @@ -772,6 +846,138 @@ class TestWorkflowService: assert draft_workflow.serialized_features == json.dumps(legacy_features) sqlite_session.refresh(draft_workflow) assert draft_workflow.serialized_features == json.dumps(legacy_features) + get_published_workflow_by_id.assert_called_once_with( + app_model=app, + workflow_id=source_workflow.id, + session=sqlite_session, + for_update=True, + ) + retire_unowned.assert_called_once_with( + tenant_id=app.tenant_id, + agent_ids={"retired-agent"}, + account_id=account.id, + ) + + def test_restore_historical_inline_agent_after_current_pointer_moves_uses_real_clone( + self, + workflow_service: WorkflowService, + sqlite_session: Session, + ) -> None: + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=None) + account = TestWorkflowAssociatedDataFactory.create_account() + graph = { + "nodes": [ + { + "id": "agent-node", + "data": { + "type": "agent", + "version": "2", + "agent_node_kind": "dify_agent", + }, + } + ], + "edges": [], + } + historical = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="historical-workflow", + version="historical-version", + graph=graph, + ) + current = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="current-workflow", + version="current-version", + graph=graph, + ) + draft = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="draft-workflow", + version=Workflow.VERSION_DRAFT, + ) + source_agent = Agent( + id="historical-agent", + tenant_id=app.tenant_id, + name="Historical inline Agent", + description="", + role="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + app_id=app.id, + workflow_id=historical.id, + workflow_node_id="agent-node", + active_config_snapshot_id="historical-snapshot", + active_config_has_model=False, + active_config_is_published=True, + status=AgentStatus.ACTIVE, + created_by=account.id, + updated_by=account.id, + ) + source_snapshot = AgentConfigSnapshot( + id="historical-snapshot", + tenant_id=app.tenant_id, + agent_id=source_agent.id, + version=1, + config_snapshot=AgentSoulConfig(config_note="historical soul"), + created_by=account.id, + ) + historical_binding = WorkflowAgentNodeBinding( + id="historical-binding", + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id=historical.id, + workflow_version=historical.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=source_agent.id, + current_snapshot_id=source_snapshot.id, + node_job_config={}, + created_by=account.id, + ) + sqlite_session.add_all([app, historical, current, draft, source_agent, source_snapshot, historical_binding]) + app.workflow_id = historical.id + sqlite_session.commit() + + app.workflow_id = current.id + sqlite_session.commit() + + WorkflowAgentRetirementService.retire_unowned( + tenant_id=app.tenant_id, + agent_ids=[source_agent.id], + account_id=account.id, + ) + sqlite_session.expire_all() + retained_agent = sqlite_session.get(Agent, source_agent.id) + assert retained_agent is not None + assert retained_agent.status is AgentStatus.ACTIVE + + with patch("services.workflow_service.app_draft_workflow_was_synced"): + restored_draft = workflow_service.restore_published_workflow_to_draft( + app_model=app, + workflow_id=historical.id, + account=account, + session=sqlite_session, + ) + + restored_binding = sqlite_session.scalar( + select(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.workflow_id == draft.id, + WorkflowAgentNodeBinding.workflow_version == Workflow.VERSION_DRAFT, + WorkflowAgentNodeBinding.node_id == "agent-node", + ) + ) + assert restored_draft is draft + assert app.workflow_id == current.id + assert sqlite_session.get(Workflow, historical.id) is historical + assert sqlite_session.get(WorkflowAgentNodeBinding, historical_binding.id) is historical_binding + assert restored_binding is not None + assert restored_binding.agent_id not in (None, source_agent.id) + assert restored_binding.current_snapshot_id not in (None, source_snapshot.id) + restored_agent = sqlite_session.get(Agent, restored_binding.agent_id) + restored_snapshot = sqlite_session.get(AgentConfigSnapshot, restored_binding.current_snapshot_id) + assert restored_agent is not None + assert restored_agent.workflow_id == draft.id + assert restored_agent.workflow_node_id == "agent-node" + assert restored_snapshot is not None + assert restored_snapshot.config_snapshot_dict == source_snapshot.config_snapshot_dict # ==================== Workflow Validation Tests ==================== # These tests verify graph structure and feature configuration validation @@ -1068,7 +1274,7 @@ class TestWorkflowService: patch("services.workflow_service.app_published_workflow_was_updated"), patch("services.workflow_service.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY), ): - result, retirement_candidates = workflow_service.publish_workflow( + result = workflow_service.publish_workflow( session=sqlite_session, app_model=app, account=account, @@ -1081,7 +1287,6 @@ class TestWorkflowService: assert result.version != Workflow.VERSION_DRAFT assert result.marked_name == "Version 1" assert result.marked_comment == "Initial release" - assert retirement_candidates == set() def test_publish_workflow_numbers_versions_from_one( self, workflow_service: WorkflowService, sqlite_session: Session @@ -1107,8 +1312,8 @@ class TestWorkflowService: DeploymentEdition.COMMUNITY, ), ): - first, _ = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - second, _ = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + first = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + second = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) assert first.version_number == 1 assert second.version_number == 2 @@ -1137,12 +1342,12 @@ class TestWorkflowService: DeploymentEdition.COMMUNITY, ), ): - published, _ = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + published = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) sqlite_session.flush() sqlite_session.delete(published) sqlite_session.flush() - republished, _ = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + republished = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) assert republished.version_number == 2 @@ -1176,7 +1381,7 @@ class TestWorkflowService: DeploymentEdition.COMMUNITY, ), ): - workflow, _ = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + workflow = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) published.append(workflow) assert [workflow.version_number for workflow in published] == [1, 1] @@ -1467,14 +1672,76 @@ class TestWorkflowService: workflow = TestWorkflowAssociatedDataFactory.create_workflow( workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id, version="v1" ) - sqlite_session.add(workflow) + inline_binding = WorkflowAgentNodeBinding( + id="inline-binding", + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_version="v1", + node_id="inline-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="inline-agent", + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + roster_binding = WorkflowAgentNodeBinding( + id="roster-binding", + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_version="v1", + node_id="roster-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="roster-agent", + current_snapshot_id="snapshot-2", + node_job_config={}, + ) + non_target_bindings = [ + WorkflowAgentNodeBinding( + id=f"non-target-{key}", + tenant_id="other-tenant" if key == "tenant" else tenant_id, + app_id="other-app" if key == "app" else app_id, + workflow_id="other-workflow" if key == "workflow" else workflow_id, + workflow_version="other-version" if key == "version" else workflow.version, + node_id=f"{key}-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=f"{key}-inline-agent", + current_snapshot_id=f"{key}-snapshot", + node_job_config={}, + ) + for key in ("tenant", "app", "workflow", "version") + ] + sqlite_session.add_all([workflow, inline_binding, roster_binding, *non_target_bindings]) sqlite_session.commit() result = workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) sqlite_session.flush() - assert result is True + assert result == ["inline-agent"] assert sqlite_session.get(Workflow, workflow_id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, inline_binding.id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, roster_binding.id) is None + for binding in non_target_bindings: + assert sqlite_session.get(WorkflowAgentNodeBinding, binding.id) is binding + + def test_delete_workflow_locks_source_until_caller_commits(self, workflow_service: WorkflowService): + workflow = TestWorkflowAssociatedDataFactory.create_workflow(version="v1") + workflow_ref = WorkflowRef( + tenant_id=workflow.tenant_id, + owner_id=workflow.app_id, + workflow_id=workflow.id, + ) + session = MagicMock(spec=Session) + session.scalar.side_effect = [workflow, None, None] + session.scalars.return_value.all.return_value = [] + + result = workflow_service.delete_workflow(session=session, workflow_ref=workflow_ref) + + stmt = session.scalar.call_args_list[0].args[0] + sql = str(stmt.compile(dialect=postgresql.dialect())) + assert result == [] + assert "FOR UPDATE" in sql + session.delete.assert_called_once_with(workflow) def test_delete_workflow_with_ref_scopes_lookup_to_app( self, workflow_service: WorkflowService, sqlite_session: Session diff --git a/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py b/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py index eb751061f17..d660ec489fe 100644 --- a/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py +++ b/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py @@ -1,8 +1,9 @@ from typing import Protocol, cast -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest +from services.agent.deletion_service import AgentDeletionService from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.workspace_service import AgentWorkspaceService from tasks.collect_agent_resources_task import ( @@ -29,6 +30,7 @@ def test_enqueue_deduplicates_ids_and_skips_empty_input(monkeypatch: pytest.Monk tenant_id="tenant-1", binding_ids=["binding-2", "binding-1", "binding-2"], workspace_ids=["workspace-1"], + purge_agent_ids=["agent-2", "", "agent-1", "agent-2"], ) delay.assert_called_once_with( @@ -36,6 +38,7 @@ def test_enqueue_deduplicates_ids_and_skips_empty_input(monkeypatch: pytest.Monk binding_ids=["binding-1", "binding-2"], workspace_ids=["workspace-1"], home_snapshot_ids=[], + purge_agent_ids=["agent-1", "agent-2"], ) @@ -56,67 +59,126 @@ def test_collection_runs_in_workspace_binding_snapshot_order(monkeypatch: pytest "collect_retired_home_snapshot", lambda **_kwargs: calls.append("home"), ) + purge = MagicMock(side_effect=lambda **_kwargs: calls.append("purge")) + monkeypatch.setattr(AgentDeletionService, "purge_archived_agents", purge) collect_agent_resources.run( tenant_id="tenant-1", workspace_ids=["workspace-1"], binding_ids=["binding-1"], home_snapshot_ids=["home-1"], + purge_agent_ids=["agent-1"], ) - assert calls == ["workspace", "binding", "home"] + assert calls == ["workspace", "binding", "home", "purge"] + purge.assert_called_once_with(tenant_id="tenant-1", agent_ids=["agent-1"]) -def test_collection_failure_propagates_and_stops_task(monkeypatch: pytest.MonkeyPatch) -> None: +def test_collection_failure_propagates_after_attempting_remaining_resources(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[str] = [] - error = RuntimeError("workspace failed") + first_error = RuntimeError("workspace-1 failed") + errors = { + "workspace-1": first_error, + "workspace-2": RuntimeError("workspace-2 failed"), + "binding-1": RuntimeError("binding-1 failed"), + "home-2": RuntimeError("home-2 failed"), + } log_exception = MagicMock() - def collect_workspace(**_kwargs: object) -> None: - calls.append("workspace") - raise error + def collect_workspace(*, workspace_id: str, **_kwargs: object) -> None: + calls.append(f"workspace:{workspace_id}") + if error := errors.get(workspace_id): + raise error + + def collect_binding(*, binding_id: str, **_kwargs: object) -> None: + calls.append(f"binding:{binding_id}") + if error := errors.get(binding_id): + raise error + + def collect_home(*, home_snapshot_id: str, **_kwargs: object) -> None: + calls.append(f"home:{home_snapshot_id}") + if error := errors.get(home_snapshot_id): + raise error monkeypatch.setattr(AgentWorkspaceService, "collect_retired_workspace", collect_workspace) - monkeypatch.setattr( - AgentWorkspaceService, - "collect_retired_binding", - lambda **_kwargs: calls.append("binding"), - ) - monkeypatch.setattr( - AgentHomeSnapshotService, - "collect_retired_home_snapshot", - lambda **_kwargs: calls.append("home"), - ) + monkeypatch.setattr(AgentWorkspaceService, "collect_retired_binding", collect_binding) + monkeypatch.setattr(AgentHomeSnapshotService, "collect_retired_home_snapshot", collect_home) monkeypatch.setattr("tasks.collect_agent_resources_task.logger.exception", log_exception) + purge = MagicMock() + monkeypatch.setattr(AgentDeletionService, "purge_archived_agents", purge) with pytest.raises(RuntimeError) as exc_info: collect_agent_resources.run( tenant_id="tenant-1", - workspace_ids=["workspace-1"], + workspace_ids=["workspace-1", "workspace-2", "workspace-3"], + binding_ids=["binding-1", "binding-2"], + home_snapshot_ids=["home-1", "home-2", "home-3"], + purge_agent_ids=["agent-1"], + ) + + assert exc_info.value.__cause__ is first_error + assert str(exc_info.value) == ( + "Failed to collect 4 retired Agent resource(s): " + "workspace:workspace-1, workspace:workspace-2, binding:binding-1, home_snapshot:home-2" + ) + assert calls == [ + "workspace:workspace-1", + "workspace:workspace-2", + "workspace:workspace-3", + "binding:binding-1", + "binding:binding-2", + "home:home-1", + "home:home-2", + "home:home-3", + ] + purge.assert_not_called() + log_exception.assert_has_calls( + [ + call( + "Failed to collect retired Agent resource", + extra={ + "tenant_id": "tenant-1", + "resource_type": resource_type, + "resource_id": resource_id, + }, + ) + for resource_type, resource_id in ( + ("workspace", "workspace-1"), + ("workspace", "workspace-2"), + ("binding", "binding-1"), + ("home_snapshot", "home-2"), + ) + ], + any_order=False, + ) + assert log_exception.call_count == 4 + + +def test_enqueue_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: + error = RuntimeError("queue unavailable") + delay = MagicMock(side_effect=error) + log_exception = MagicMock() + monkeypatch.setattr(collect_agent_resources, "delay", delay) + monkeypatch.setattr("tasks.collect_agent_resources_task.logger.exception", log_exception) + + with pytest.raises(RuntimeError) as exc_info: + enqueue_agent_resource_collection( + tenant_id="tenant-1", binding_ids=["binding-1"], + workspace_ids=["workspace-1"], home_snapshot_ids=["home-1"], + purge_agent_ids=["agent-2", "agent-1", "agent-2"], ) assert exc_info.value is error - assert calls == ["workspace"] + payload = { + "binding_ids": ["binding-1"], + "workspace_ids": ["workspace-1"], + "home_snapshot_ids": ["home-1"], + "purge_agent_ids": ["agent-1", "agent-2"], + } + delay.assert_called_once_with(tenant_id="tenant-1", **payload) log_exception.assert_called_once_with( - "Failed to collect retired Agent resource", - extra={ - "tenant_id": "tenant-1", - "resource_type": "workspace", - "resource_id": "workspace-1", - }, - ) - - -def test_enqueue_failure_is_best_effort(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - collect_agent_resources, - "delay", - MagicMock(side_effect=RuntimeError("queue unavailable")), - ) - - enqueue_agent_resource_collection( - tenant_id="tenant-1", - binding_ids=["binding-1"], + "Failed to enqueue retired Agent resource collection", + extra={"tenant_id": "tenant-1", **payload}, ) diff --git a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py index 3f0ce2475ef..125e447a10e 100644 --- a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py +++ b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py @@ -6,9 +6,12 @@ from uuid import uuid4 import pytest from sqlalchemy.orm import Session +import tasks.remove_app_and_related_data_task as remove_app_task_module +from enums import DeploymentEdition from graphon.enums import WorkflowExecutionStatus from libs.archive_storage import ArchiveStorageNotConfiguredError from models import AppStar +from models.agent import WorkflowAgentBindingType, WorkflowAgentNodeBinding from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.workflow import WorkflowArchiveLog from tasks.remove_app_and_related_data_task import ( @@ -17,10 +20,104 @@ from tasks.remove_app_and_related_data_task import ( _delete_archived_workflow_run_files, _delete_draft_variable_offload_data, _delete_draft_variables, + _delete_workflow_agent_node_bindings, delete_draft_variables_batch, ) +def test_delete_workflow_agent_node_bindings_is_scoped_to_tenant_and_app(sqlite_session: Session) -> None: + target = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version="draft", + node_id="node-1", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="agent-1", + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + kept = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-2", + workflow_id="workflow-2", + workflow_version="draft", + node_id="node-2", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="agent-2", + current_snapshot_id="snapshot-2", + node_job_config={}, + ) + other_tenant = WorkflowAgentNodeBinding( + tenant_id="tenant-2", + app_id="app-1", + workflow_id="workflow-3", + workflow_version="draft", + node_id="node-3", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="agent-3", + current_snapshot_id="snapshot-3", + node_job_config={}, + ) + sqlite_session.add_all([target, kept, other_tenant]) + sqlite_session.commit() + target_id = target.id + kept_id = kept.id + other_tenant_id = other_tenant.id + + _delete_workflow_agent_node_bindings("tenant-1", "app-1") + + sqlite_session.expire_all() + assert sqlite_session.get(WorkflowAgentNodeBinding, target_id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, kept_id) is not None + assert sqlite_session.get(WorkflowAgentNodeBinding, other_tenant_id) is not None + + +def test_app_cleanup_removes_agent_bindings_before_workflows(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + monkeypatch.setattr(remove_app_task_module.dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY) + other_cleanup_names = ( + "_delete_app_model_configs", + "_delete_app_site", + "_delete_app_mcp_servers", + "_delete_app_api_tokens", + "_delete_installed_apps", + "_delete_app_stars", + "_delete_recommended_apps", + "_delete_app_annotation_data", + "_delete_app_dataset_joins", + "_delete_app_workflow_runs", + "_delete_app_workflow_node_executions", + "_delete_app_workflow_app_logs", + "_delete_app_conversations", + "_delete_app_messages", + "_delete_workflow_tool_providers", + "_delete_app_tag_bindings", + "_delete_end_users", + "_delete_trace_app_configs", + "_delete_conversation_variables", + "_delete_draft_variables", + "_delete_app_triggers", + "_delete_workflow_plugin_triggers", + "_delete_workflow_webhook_triggers", + "_delete_workflow_schedule_plans", + "_delete_workflow_trigger_logs", + ) + for name in other_cleanup_names: + monkeypatch.setattr(remove_app_task_module, name, MagicMock()) + + delete_bindings = MagicMock(side_effect=lambda *_args: events.append("bindings")) + delete_workflows = MagicMock(side_effect=lambda *_args: events.append("workflows")) + monkeypatch.setattr(remove_app_task_module, "_delete_workflow_agent_node_bindings", delete_bindings) + monkeypatch.setattr(remove_app_task_module, "_delete_app_workflows", delete_workflows) + + remove_app_task_module.remove_app_and_related_data_task.run(tenant_id="tenant-1", app_id="app-1") + + assert events == ["bindings", "workflows"] + delete_bindings.assert_called_once_with("tenant-1", "app-1") + delete_workflows.assert_called_once_with("tenant-1", "app-1") + + class TestDeleteDraftVariablesBatch: def test_delete_draft_variables_batch_invalid_batch_size(self): """Test that invalid batch size raises ValueError.""" diff --git a/api/tests/unit_tests/test_constants.py b/api/tests/unit_tests/test_constants.py new file mode 100644 index 00000000000..e40744a894a --- /dev/null +++ b/api/tests/unit_tests/test_constants.py @@ -0,0 +1,25 @@ +import importlib + +import pytest + +import constants +from configs import dify_config + + +@pytest.mark.parametrize("etl_type", ["SelfHosted", "Unstructured"]) +def test_document_extensions_include_odt_for_document_etl_modes(monkeypatch: pytest.MonkeyPatch, etl_type: str) -> None: + original_etl_type = dify_config.ETL_TYPE + original_unstructured_api_url = dify_config.UNSTRUCTURED_API_URL + + try: + monkeypatch.setattr(dify_config, "ETL_TYPE", etl_type) + monkeypatch.setattr(dify_config, "UNSTRUCTURED_API_URL", None) + + reloaded_constants = importlib.reload(constants) + + assert "odt" in reloaded_constants.DOCUMENT_EXTENSIONS + assert "ODT" in reloaded_constants.DOCUMENT_EXTENSIONS + finally: + monkeypatch.setattr(dify_config, "ETL_TYPE", original_etl_type) + monkeypatch.setattr(dify_config, "UNSTRUCTURED_API_URL", original_unstructured_api_url) + importlib.reload(constants) diff --git a/api/uv.lock b/api/uv.lock index 76d41e89ef8..4616ba1845c 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -613,16 +613,16 @@ wheels = [ [[package]] name = "boto3" -version = "1.43.56" +version = "1.43.71" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore" }, { name = "jmespath" }, { name = "s3transfer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/53/05/23e1aa8c9e4b0399a61e7fd65c4f9cc0625121f24760e37471f776404abb/boto3-1.43.56.tar.gz", hash = "sha256:57c90df9fb026f2e6ae22530861198130203733c5c9ec4e5cca3a4037f5a8db4", size = 112673, upload-time = "2026-07-24T19:31:48.606Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/ed/3220ca6396a7ce00a49aa6f43f20cae05ec6ac386e2b5d668adfad0b3eed/boto3-1.43.71.tar.gz", hash = "sha256:e3eddb6346ee23c895dd98e5c2cfc72cb5e5ed111aa6d4331f9e903a3e3bf4ee", size = 112630, upload-time = "2026-08-13T19:20:35.332Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b8/57/3a960c9f581c00f2a591901b46e035ff79ab3956d16607f12306b3b8d483/boto3-1.43.56-py3-none-any.whl", hash = "sha256:feb699d4ab241ef5c1b80bb58277be2aaad365cd4b672d7817e0bc59ee45131b", size = 140026, upload-time = "2026-07-24T19:31:47.155Z" }, + { url = "https://files.pythonhosted.org/packages/4f/16/7b54e7f5271756ab13beaaf6da03f62edcc3ee5ded651c2b17297e107040/boto3-1.43.71-py3-none-any.whl", hash = "sha256:3bb46661dc33121e56e3706753069d18f8dbcb644599cf098206546338602e16", size = 140024, upload-time = "2026-08-13T19:20:32.542Z" }, ] [[package]] @@ -645,16 +645,16 @@ bedrock-runtime = [ [[package]] name = "botocore" -version = "1.43.56" +version = "1.43.72" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jmespath" }, { name = "python-dateutil" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b9/cc/7f84a5d3071fe878380e9f610ab36ca87b8cbbc4aa81ba2727f90e1f3ea3/botocore-1.43.56.tar.gz", hash = "sha256:6c01f85f0ff9863076f4c761e74ee3aa96c5ccc1ad09fc1efd62ef8f2d22bf57", size = 15733117, upload-time = "2026-07-24T19:31:38.125Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a1/99/a8cfeaea98d5085a493af909d09d174466482235a7fda291be18c9a5a76e/botocore-1.43.72.tar.gz", hash = "sha256:1b878c69081e8e9d55aa4c0d85683e7b07f0e274a5554662f9507a46641be3d2", size = 15949280, upload-time = "2026-08-14T19:24:48.841Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5c/cd/86fe9e659e9699f62f8dd5ecd8c6725474334b23cab8aa71d82b5f56f1a4/botocore-1.43.56-py3-none-any.whl", hash = "sha256:aafc741f1b10f6fd63253eaf6ea029680c1ff436d87e1b8969d62aefa0c76976", size = 15418773, upload-time = "2026-07-24T19:31:34.758Z" }, + { url = "https://files.pythonhosted.org/packages/f4/27/35814fd8701a6b8be0aa47a7fcbd1ea0706189410a47d71ff29ddcc3ad4d/botocore-1.43.72-py3-none-any.whl", hash = "sha256:de5a1bcf8d7602c6cefc15016f15dad82981e339192531f31fa9483e11feea47", size = 15641880, upload-time = "2026-08-14T19:24:45.882Z" }, ] [[package]] @@ -1642,7 +1642,7 @@ requires-dist = [ { name = "aliyun-log-python-sdk", specifier = "==0.9.44" }, { name = "azure-identity", specifier = ">=1.25.3,<2.0.0" }, { name = "bleach", specifier = ">=6.4.0,<7.0.0" }, - { name = "boto3", specifier = ">=1.43.56,<2.0.0" }, + { name = "boto3", specifier = ">=1.43.71,<2.0.0" }, { name = "celery", specifier = ">=5.6.3,<6.0.0" }, { name = "croniter", specifier = ">=6.2.2,<7.0.0" }, { name = "dify-agent", editable = "../dify-agent" }, @@ -1658,7 +1658,7 @@ requires-dist = [ { name = "gevent-websocket", specifier = "==0.10.1" }, { name = "gmpy2", specifier = ">=2.3.0,<3.0.0" }, { name = "google-api-python-client", specifier = ">=2.198.0,<3.0.0" }, - { name = "google-cloud-aiplatform", specifier = ">=1.160.0,<2.0.0" }, + { name = "google-cloud-aiplatform", specifier = ">=1.164.0,<2.0.0" }, { name = "graphon", specifier = "==0.7.0" }, { name = "gunicorn", specifier = ">=26.0.0,<27.0.0" }, { name = "httpx", extras = ["socks"], specifier = "==0.28.1" }, @@ -1755,7 +1755,7 @@ storage = [ { name = "bce-python-sdk", specifier = "==0.9.76" }, { name = "cos-python-sdk-v5", specifier = ">=1.9.44,<2.0.0" }, { name = "esdk-obs-python", specifier = ">=3.26.6,<4.0.0" }, - { name = "google-cloud-storage", specifier = ">=3.13.0,<4.0.0" }, + { name = "google-cloud-storage", specifier = ">=3.13.1,<4.0.0" }, { name = "opendal", specifier = "==0.46.0" }, { name = "oss2", specifier = ">=2.19.1,<3.0.0" }, { name = "supabase", specifier = ">=2.31.0,<3.0.0" }, @@ -2844,7 +2844,7 @@ wheels = [ [[package]] name = "google-cloud-aiplatform" -version = "1.160.0" +version = "1.164.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "certifi" }, @@ -2861,9 +2861,9 @@ dependencies = [ { name = "pydantic" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d6/c5/dad5053ce2bbf53079274c4781f7bdf45d1f85bfe0ea8fad88cd39fad52d/google_cloud_aiplatform-1.160.0.tar.gz", hash = "sha256:186a8db5099eda0e3cd3ecc73a4716d48c82fa1f00501582eacd541a6aa60534", size = 11174223, upload-time = "2026-07-08T00:48:13.036Z" } +sdist = { url = "https://files.pythonhosted.org/packages/81/bd/ddc72de71fd76ad40e88d23a634d0e6f2108330b98f349463016e5f7c0c1/google_cloud_aiplatform-1.164.0.tar.gz", hash = "sha256:51fb881523631e591a3d11a46cb9d7bd390e05127fe1d08c596ddce3e45fed2d", size = 11298864, upload-time = "2026-08-12T20:40:52.538Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/32/06/230752e7697ba83f92a0d124f42b43e170ecf8c212a97a2f2fc8c788d4c1/google_cloud_aiplatform-1.160.0-py2.py3-none-any.whl", hash = "sha256:4886e035b5baad3fe52adcec1c9a9f8312b5c42d7dfed8e13b420d59e6f3b82a", size = 9369771, upload-time = "2026-07-08T00:48:09.327Z" }, + { url = "https://files.pythonhosted.org/packages/4a/29/8bb9f8df34e925d4ec895104e2aa33fea5ac9cf43dd96654888224da47ac/google_cloud_aiplatform-1.164.0-py2.py3-none-any.whl", hash = "sha256:042dbbd9f48524ad4ac34c25b00be4f4ca8b1b38fd6d5d277cc80501777974a0", size = 9432429, upload-time = "2026-08-12T20:40:48.936Z" }, ] [[package]] @@ -2916,7 +2916,7 @@ wheels = [ [[package]] name = "google-cloud-storage" -version = "3.13.0" +version = "3.13.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "google-api-core" }, @@ -2926,9 +2926,9 @@ dependencies = [ { name = "google-resumable-media" }, { name = "requests" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e3/25/355ed97c1723c787dfaa888808d55db18371f82c38ff862357b1e902cd19/google_cloud_storage-3.13.0.tar.gz", hash = "sha256:d11d8706ea1520fba0f21043bcb7897caf7015d76ce1ad9a4f60237e4d7a9f6c", size = 17340960, upload-time = "2026-07-13T19:10:07.524Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ce/7e/73bb7512df1d1aad6ce3f9aed847cd40e0cd400ba4a85d86ab8eb412e9cc/google_cloud_storage-3.13.1.tar.gz", hash = "sha256:a80bf8cac2794808aa61c50c5f769ecbbe2d10331bacd0d69d30e59b14b346b2", size = 17341051, upload-time = "2026-08-06T06:24:42.229Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/81/e8/b3678a0931ee7d4b3fdaf0813e6206d66e0922b2c26d912f308728b5b95a/google_cloud_storage-3.13.0-py3-none-any.whl", hash = "sha256:648af3ef8a6acc674e1359d3c920c67eb89a7a5ab66b336bd3ac43fed6b5ab84", size = 341428, upload-time = "2026-07-13T19:09:52.39Z" }, + { url = "https://files.pythonhosted.org/packages/06/6f/d69f0e185e08ddb58c323a0a935af2b492907b5de362bc08933b0a3b5644/google_cloud_storage-3.13.1-py3-none-any.whl", hash = "sha256:98208de6c21e85cecd3eb44551894efff33d98365500e178867d4305854a770a", size = 341486, upload-time = "2026-08-06T06:23:36.548Z" }, ] [[package]] diff --git a/dify-agent-runtime/README.md b/dify-agent-runtime/README.md index 59b1d451167..00b976c7742 100644 --- a/dify-agent-runtime/README.md +++ b/dify-agent-runtime/README.md @@ -14,6 +14,21 @@ cmd/ internal/ - internal implementations ``` +## Job execution modes + +`POST /v1/jobs/run` accepts an optional `mode` field: + +- `pty` (default) keeps the interactive tmux PTY path. stdout and stderr are + merged, sanitized, and written to `output.log`; the job accepts `/input`. +- `stdio` keeps tmux as the lifecycle owner but gives the child `/dev/null` as + stdin and captures stdout and stderr through separate pipes. Public output + and pagination read stdout from `output.log`; private diagnostics are written + to `stderr.log`. A stdio job completes only after both streams reach EOF and + does not accept `/input`. + +The response models are identical in both modes. Use `stdio` for bounded, +machine-readable control commands and `pty` for interactive jobs. + ## Building ```bash @@ -59,12 +74,12 @@ Each agent job runs inside a Landlock sandbox that restricts filesystem access: | Access | Paths (defaults) | | -------------------- | ------------------------------------------------------------------------------------------------------------- | -| **Read-Write** | `$HOME` (always, includes `$CWD/.tmp` as `TMPDIR`) | +| **Read-Write** | `$HOME` and the job's `cwd` (also used directly as `TMPDIR`, `TMP`, and `TEMP`) | | **Read-Write (dev)** | `/dev/null`, `/dev/zero`, `/dev/urandom`, `/dev/random`, `/dev/tty` | | **Read-Only + Exec** | `/usr`, `/bin`, `/sbin`, `/lib`, `/lib64`, `/etc`, `/proc`, `/opt/dify-agent-tools`, `/opt/homebrew`, `/snap` | | **Denied** | Everything else (`/tmp`, other agents' homes, `/var`, `/srv`, etc.) | -The runner automatically creates `$CWD/.tmp` and sets `TMPDIR`, `TMP`, `TEMP` to it, so temp files stay isolated per workspace. +The runner sets `TMPDIR`, `TMP`, and `TEMP` directly to the job's `cwd`. It does not create a separate temp directory, so the active Workspace is both the working directory and temp space. ### Environment Variables diff --git a/dify-agent-runtime/cmd/runner/main.go b/dify-agent-runtime/cmd/runner/main.go index 29f35b4867d..9d9a669f885 100644 --- a/dify-agent-runtime/cmd/runner/main.go +++ b/dify-agent-runtime/cmd/runner/main.go @@ -1,9 +1,9 @@ // shellctl-runner is the Go replacement for the previously generated bash+python // runner script. It is invoked by tmux as: // -// shellctl-runner +// shellctl-runner [pty|stdio] // -// The binary operates in two modes: +// The binary operates in two process roles: // // 1. Parent mode (default): waits for start-gate, loads env, forks child, // waits for exit, writes exit artifacts. @@ -17,16 +17,19 @@ package main import ( "encoding/json" "fmt" + "io" "os" "os/exec" "os/signal" "path/filepath" "strings" + "sync" "syscall" "time" "github.com/langgenius/dify/dify-agent-runtime/internal/cmdutil" "github.com/langgenius/dify/dify-agent-runtime/internal/envvar" + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" "github.com/langgenius/dify/dify-agent-runtime/internal/landlock" ) @@ -39,21 +42,25 @@ func main() { } // parentMode is the entry point when called by tmux. -// Args: shellctl-runner +// Args: shellctl-runner [pty|stdio] func parentMode() { if len(os.Args) < 4 { - cmdutil.HandleError(fmt.Errorf("bad args"), 125, "usage: shellctl-runner ") + cmdutil.HandleError(fmt.Errorf("bad args"), 125, "usage: shellctl-runner [pty|stdio]") } jobDir := os.Args[1] // jobID := os.Args[2] // unused in parent but passed for compat cwd := os.Args[3] + modeRaw := "" + if len(os.Args) >= 5 { + modeRaw = os.Args[4] + } + mode, err := jobmode.Parse(modeRaw) + cmdutil.HandleError(err, 125, "parse job mode") scriptPath := filepath.Join(jobDir, "script") envPath := filepath.Join(jobDir, ".job-env.json") startGate := filepath.Join(jobDir, "start-gate") - exitCodePath := filepath.Join(jobDir, "runner-exit-code") - endedAtPath := filepath.Join(jobDir, "runner-ended-at") // Wait for start-gate. for { @@ -77,6 +84,11 @@ func parentMode() { envOverlay := loadEnvJSON(envPath) env = mergeEnv(env, envOverlay) + env = mergeEnv(env, map[string]string{ + "TMPDIR": cwd, + "TMP": cwd, + "TEMP": cwd, + }) // Ensure HOME exists. home := envGet(env, "HOME") @@ -84,14 +96,6 @@ func parentMode() { cmdutil.HandleError(os.MkdirAll(home, 0755), 125, "mkdir HOME %s", home) } - // Create a per-workspace temp directory under cwd and inject TMPDIR. - // This avoids granting RW access to the shared /tmp. - agentTmp := filepath.Join(cwd, ".tmp") - cmdutil.HandleError(os.MkdirAll(agentTmp, 0755), 125, "mkdir TMPDIR %s", agentTmp) - env = setEnvIfEmpty(env, "TMPDIR", agentTmp) - env = setEnvIfEmpty(env, "TMP", agentTmp) - env = setEnvIfEmpty(env, "TEMP", agentTmp) - // Determine if path isolation is enabled. enableIsolation := envvar.PathIsolationEnabled() @@ -104,9 +108,6 @@ func parentMode() { cmd := exec.Command(self, childArgs...) cmd.Env = env - cmd.Stdin = os.Stdin - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr cmd.Dir = cwd // Forward signals to child. @@ -120,23 +121,154 @@ func parentMode() { } }() - err := cmd.Run() + exitCode := runCommandAndRecordExit(cmd, jobDir, mode) + os.Exit(exitCode) +} - exitCode := 0 - if err != nil { - if exitErr, ok := err.(*exec.ExitError); ok { - exitCode = exitErr.ExitCode() - } else { - exitCode = 125 - } +// runCommandAndRecordExit publishes the existing runner artifacts after the +// child path returns. In stdio mode that return includes both stream drains; +// PTY mode still relies on its separate pipe-drain finalizer. +func runCommandAndRecordExit(cmd *exec.Cmd, jobDir string, mode jobmode.Mode) int { + var exitCode int + if mode == jobmode.Stdio { + exitCode = runStdio(cmd, jobDir) + } else { + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + exitCode = runPTY(cmd) } endedAt := time.Now().UTC().Format("2006-01-02T15:04:05Z") - writeAtomic(exitCodePath, fmt.Sprintf("%d", exitCode)) - writeAtomic(endedAtPath, endedAt) + writeAtomic(filepath.Join(jobDir, "runner-exit-code"), fmt.Sprintf("%d", exitCode)) + writeAtomic(filepath.Join(jobDir, "runner-ended-at"), endedAt) + return exitCode +} - os.Exit(exitCode) +func runPTY(cmd *exec.Cmd) int { + return commandExitCode(cmd.Run()) +} + +// runStdio captures stdout and stderr independently and does not return until +// both streams reach EOF and their files are closed. +func runStdio(cmd *exec.Cmd, jobDir string) int { + outputFile, err := openCaptureFile(filepath.Join(jobDir, "output.log")) + if err != nil { + return runnerError("open stdout capture", err) + } + stderrFile, err := openCaptureFile(filepath.Join(jobDir, "stderr.log")) + if err != nil { + _ = outputFile.Close() + return runnerError("open stderr capture", err) + } + + stdoutReader, stdoutWriter, err := os.Pipe() + if err != nil { + _ = outputFile.Close() + _ = stderrFile.Close() + return runnerError("create stdout pipe", err) + } + stderrReader, stderrWriter, err := os.Pipe() + if err != nil { + _ = stdoutReader.Close() + _ = stdoutWriter.Close() + _ = outputFile.Close() + _ = stderrFile.Close() + return runnerError("create stderr pipe", err) + } + stdinFile, err := os.Open(os.DevNull) + if err != nil { + _ = stdoutReader.Close() + _ = stdoutWriter.Close() + _ = stderrReader.Close() + _ = stderrWriter.Close() + _ = outputFile.Close() + _ = stderrFile.Close() + return runnerError("open stdin", err) + } + + cmd.Stdin = stdinFile + cmd.Stdout = stdoutWriter + cmd.Stderr = stderrWriter + + var captureErrors []error + var captureErrorsMu sync.Mutex + recordError := func(operation string, err error) { + if err != nil { + captureErrorsMu.Lock() + captureErrors = append(captureErrors, fmt.Errorf("%s: %w", operation, err)) + captureErrorsMu.Unlock() + } + } + + var drains sync.WaitGroup + drains.Add(2) + go func() { + defer drains.Done() + _, copyErr := io.Copy(outputFile, stdoutReader) + recordError("copy stdout", copyErr) + recordError("close stdout reader", stdoutReader.Close()) + }() + go func() { + defer drains.Done() + _, copyErr := io.Copy(stderrFile, stderrReader) + recordError("copy stderr", copyErr) + recordError("close stderr reader", stderrReader.Close()) + }() + + startErr := cmd.Start() + recordError("close stdout writer", stdoutWriter.Close()) + recordError("close stderr writer", stderrWriter.Close()) + recordError("close stdin", stdinFile.Close()) + + var waitErr error + if startErr != nil { + recordError("start child", startErr) + } else { + waitErr = cmd.Wait() + } + + // Descendants may retain either write end after the direct child exits. In + // that case the runner intentionally remains alive until both reach EOF. + drains.Wait() + recordError("close stdout capture", outputFile.Close()) + recordError("close stderr capture", stderrFile.Close()) + + for _, captureErr := range captureErrors { + fmt.Fprintf(os.Stderr, "shellctl-runner: %v\n", captureErr) + } + if len(captureErrors) > 0 { + return 125 + } + return commandExitCode(waitErr) +} + +func openCaptureFile(path string) (*os.File, error) { + file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600) + if err != nil { + return nil, err + } + if err := file.Chmod(0600); err != nil { + _ = file.Close() + return nil, err + } + return file, nil +} + +func runnerError(operation string, err error) int { + fmt.Fprintf(os.Stderr, "shellctl-runner: %s: %v\n", operation, err) + return 125 +} + +func commandExitCode(err error) int { + if err == nil { + return 0 + } + if exitErr, ok := err.(*exec.ExitError); ok { + return exitErr.ExitCode() + } + return 125 } // childMode applies Landlock (if --landlock flag) and exec's the user script. @@ -264,14 +396,6 @@ func envGet(env []string, key string) string { return "" } -// setEnvIfEmpty sets key=value in the env slice only if the key is not already present. -func setEnvIfEmpty(env []string, key, value string) []string { - if envGet(env, key) != "" { - return env - } - return append(env, key+"="+value) -} - // writeAtomic writes value to dest via a temp file + rename. func writeAtomic(dest, value string) { tmp := fmt.Sprintf("%s.tmp.%d", dest, os.Getpid()) diff --git a/dify-agent-runtime/cmd/runner/main_test.go b/dify-agent-runtime/cmd/runner/main_test.go new file mode 100644 index 00000000000..9bf82100a2c --- /dev/null +++ b/dify-agent-runtime/cmd/runner/main_test.go @@ -0,0 +1,294 @@ +package main + +import ( + "bytes" + "os" + "os/exec" + "path/filepath" + "strconv" + "syscall" + "testing" + "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" +) + +func TestRunStdioCapturesCompleteSeparatedStreams(t *testing.T) { + jobDir := t.TempDir() + cmd := exec.Command(os.Args[0], "-test.run=^TestStdioHelperProcess$") + cmd.Env = append(os.Environ(), "SHELLCTL_STDIO_HELPER=large-output") + + if exitCode := runStdio(cmd, jobDir); exitCode != 0 { + t.Fatalf("runStdio exit code = %d, want 0", exitCode) + } + + wantStdout := bytes.Repeat([]byte("stdout-payload\n"), 16*1024) + wantStderr := bytes.Repeat([]byte("stderr-payload\n"), 16*1024) + gotStdout, err := os.ReadFile(filepath.Join(jobDir, "output.log")) + if err != nil { + t.Fatalf("read output.log: %v", err) + } + gotStderr, err := os.ReadFile(filepath.Join(jobDir, "stderr.log")) + if err != nil { + t.Fatalf("read stderr.log: %v", err) + } + if !bytes.Equal(gotStdout, wantStdout) { + t.Errorf("stdout capture length = %d, want %d", len(gotStdout), len(wantStdout)) + } + if !bytes.Equal(gotStderr, wantStderr) { + t.Errorf("stderr capture length = %d, want %d", len(gotStderr), len(wantStderr)) + } + for _, name := range []string{"output.log", "stderr.log"} { + info, err := os.Stat(filepath.Join(jobDir, name)) + if err != nil { + t.Fatalf("stat %s: %v", name, err) + } + if got := info.Mode().Perm(); got != 0600 { + t.Errorf("%s permissions = %#o, want 0600", name, got) + } + } +} + +func TestRunStdioUsesNonTTYStreams(t *testing.T) { + jobDir := t.TempDir() + cmd := exec.Command("sh", "-c", `if [ -t 0 ] || [ -t 1 ] || [ -t 2 ]; then exit 1; fi; printf 'stdout-only'; printf 'warning' >&2`) + + if exitCode := runStdio(cmd, jobDir); exitCode != 0 { + t.Fatalf("runStdio exit code = %d, want 0", exitCode) + } + stdout, err := os.ReadFile(filepath.Join(jobDir, "output.log")) + if err != nil { + t.Fatal(err) + } + stderr, err := os.ReadFile(filepath.Join(jobDir, "stderr.log")) + if err != nil { + t.Fatal(err) + } + if string(stdout) != "stdout-only" { + t.Errorf("stdout = %q, want stdout-only", stdout) + } + if string(stderr) != "warning" { + t.Errorf("stderr = %q, want warning", stderr) + } +} + +func TestRunStdioWaitsForBothDescendantStreamsBeforePublishingExit(t *testing.T) { + jobDir := t.TempDir() + cmd := exec.Command(os.Args[0], "-test.run=^TestStdioHelperProcess$") + cmd.Env = mergeEnv(os.Environ(), map[string]string{ + "SHELLCTL_STDIO_HELPER": "spawn-descendants", + "SHELLCTL_STDIO_JOB_DIR": jobDir, + }) + done := make(chan struct{}) + var exitCode int + go func() { + exitCode = runCommandAndRecordExit(cmd, jobDir, jobmode.Stdio) + close(done) + }() + + stdoutRelease := filepath.Join(jobDir, "release-stdout") + stderrRelease := filepath.Join(jobDir, "release-stderr") + t.Cleanup(func() { + _ = os.WriteFile(stdoutRelease, nil, 0600) + _ = os.WriteFile(stderrRelease, nil, 0600) + for _, name := range []string{"stdout-closed", "stderr-closed"} { + if !waitForPath(filepath.Join(jobDir, name), 5*time.Second) { + t.Errorf("cleanup timed out waiting for %s", name) + } + } + select { + case <-done: + case <-time.After(5 * time.Second): + t.Error("cleanup timed out waiting for runner completion") + } + }) + + waitForTestFile(t, filepath.Join(jobDir, "direct-child-exited")) + assertRunnerRemainsIncomplete(t, done) + assertExitArtifactsAbsent(t, jobDir) + + if err := os.WriteFile(stdoutRelease, nil, 0600); err != nil { + t.Fatal(err) + } + waitForTestFile(t, filepath.Join(jobDir, "stdout-closed")) + assertRunnerRemainsIncomplete(t, done) + assertExitArtifactsAbsent(t, jobDir) + + if err := os.WriteFile(stderrRelease, nil, 0600); err != nil { + t.Fatal(err) + } + waitForTestFile(t, filepath.Join(jobDir, "stderr-closed")) + + select { + case <-done: + if exitCode != 0 { + t.Fatalf("exit code = %d, want 0", exitCode) + } + case <-time.After(5 * time.Second): + t.Fatal("runner did not complete after both descendant streams reached EOF") + } + waitForTestFile(t, filepath.Join(jobDir, "runner-exit-code")) + waitForTestFile(t, filepath.Join(jobDir, "runner-ended-at")) + + stdout, err := os.ReadFile(filepath.Join(jobDir, "output.log")) + if err != nil { + t.Fatal(err) + } + stderr, err := os.ReadFile(filepath.Join(jobDir, "stderr.log")) + if err != nil { + t.Fatal(err) + } + if string(stdout) != "stdout-tail" { + t.Errorf("stdout = %q, want stdout-tail", stdout) + } + if string(stderr) != "stderr-tail" { + t.Errorf("stderr = %q, want stderr-tail", stderr) + } +} + +func TestRunStdioPreservesNonZeroExitCode(t *testing.T) { + jobDir := t.TempDir() + cmd := exec.Command("sh", "-c", "exit 23") + + if exitCode := runCommandAndRecordExit(cmd, jobDir, jobmode.Stdio); exitCode != 23 { + t.Fatalf("exit code = %d, want 23", exitCode) + } + exitCodeArtifact, err := os.ReadFile(filepath.Join(jobDir, "runner-exit-code")) + if err != nil { + t.Fatal(err) + } + if string(exitCodeArtifact) != "23\n" { + t.Errorf("runner-exit-code = %q, want 23", exitCodeArtifact) + } +} + +func TestStdioHelperProcess(t *testing.T) { + switch os.Getenv("SHELLCTL_STDIO_HELPER") { + case "": + return + case "large-output": + stdout := bytes.Repeat([]byte("stdout-payload\n"), 16*1024) + stderr := bytes.Repeat([]byte("stderr-payload\n"), 16*1024) + _, _ = os.Stdout.Write(stdout) + _, _ = os.Stderr.Write(stderr) + os.Exit(0) + case "spawn-descendants": + spawnStdioDescendants() + case "hold-stdout", "hold-stderr": + holdStdioStream(os.Getenv("SHELLCTL_STDIO_HELPER")) + default: + os.Exit(125) + } +} + +func spawnStdioDescendants() { + jobDir := os.Getenv("SHELLCTL_STDIO_JOB_DIR") + parentPID := strconv.Itoa(os.Getpid()) + stdoutHolder := exec.Command(os.Args[0], "-test.run=^TestStdioHelperProcess$") + stdoutHolder.Env = mergeEnv(os.Environ(), map[string]string{ + "SHELLCTL_STDIO_HELPER": "hold-stdout", + "SHELLCTL_STDIO_JOB_DIR": jobDir, + "SHELLCTL_STDIO_PARENT_PID": parentPID, + }) + stdoutHolder.Stdout = os.Stdout + if err := stdoutHolder.Start(); err != nil { + os.Exit(125) + } + + stderrHolder := exec.Command(os.Args[0], "-test.run=^TestStdioHelperProcess$") + stderrHolder.Env = mergeEnv(os.Environ(), map[string]string{ + "SHELLCTL_STDIO_HELPER": "hold-stderr", + "SHELLCTL_STDIO_JOB_DIR": jobDir, + }) + stderrHolder.Stderr = os.Stderr + if err := stderrHolder.Start(); err != nil { + os.Exit(125) + } + + if !waitForPath(filepath.Join(jobDir, "stdout-ready"), 5*time.Second) || + !waitForPath(filepath.Join(jobDir, "stderr-ready"), 5*time.Second) { + os.Exit(125) + } + os.Exit(0) +} + +func holdStdioStream(mode string) { + jobDir := os.Getenv("SHELLCTL_STDIO_JOB_DIR") + streamName := mode[len("hold-"):] + if err := os.WriteFile(filepath.Join(jobDir, streamName+"-ready"), nil, 0600); err != nil { + os.Exit(125) + } + if mode == "hold-stdout" { + parentPID, err := strconv.Atoi(os.Getenv("SHELLCTL_STDIO_PARENT_PID")) + if err != nil || !waitForProcessExit(parentPID, 5*time.Second) { + os.Exit(125) + } + if err := os.WriteFile(filepath.Join(jobDir, "direct-child-exited"), nil, 0600); err != nil { + os.Exit(125) + } + } + + if !waitForPath(filepath.Join(jobDir, "release-"+streamName), 5*time.Second) { + os.Exit(125) + } + if mode == "hold-stdout" { + _, _ = os.Stdout.WriteString("stdout-tail") + _ = os.Stdout.Close() + } else { + _, _ = os.Stderr.WriteString("stderr-tail") + _ = os.Stderr.Close() + } + if err := os.WriteFile(filepath.Join(jobDir, streamName+"-closed"), nil, 0600); err != nil { + os.Exit(125) + } + os.Exit(0) +} + +func waitForProcessExit(pid int, timeout time.Duration) bool { + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if err := syscall.Kill(pid, 0); err == syscall.ESRCH { + return true + } + time.Sleep(5 * time.Millisecond) + } + return false +} + +func waitForPath(path string, timeout time.Duration) bool { + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if _, err := os.Stat(path); err == nil { + return true + } + time.Sleep(5 * time.Millisecond) + } + return false +} + +func waitForTestFile(t *testing.T, path string) { + t.Helper() + if !waitForPath(path, 5*time.Second) { + t.Fatalf("timed out waiting for %s", filepath.Base(path)) + } +} + +func assertRunnerRemainsIncomplete(t *testing.T, done <-chan struct{}) { + t.Helper() + timer := time.NewTimer(100 * time.Millisecond) + defer timer.Stop() + select { + case <-done: + t.Fatal("runner completed before both streams reached EOF") + case <-timer.C: + } +} + +func assertExitArtifactsAbsent(t *testing.T, jobDir string) { + t.Helper() + for _, name := range []string{"runner-exit-code", "runner-ended-at"} { + if _, err := os.Stat(filepath.Join(jobDir, name)); !os.IsNotExist(err) { + t.Fatalf("%s became visible before both streams reached EOF: %v", name, err) + } + } +} diff --git a/dify-agent-runtime/docker/Dockerfile b/dify-agent-runtime/docker/Dockerfile index b0ed2e499fa..f8d38413e94 100644 --- a/dify-agent-runtime/docker/Dockerfile +++ b/dify-agent-runtime/docker/Dockerfile @@ -71,8 +71,9 @@ COPY --from=go-builder /bin/shellctl-runner /usr/local/bin/shellctl-runner COPY --from=go-builder /bin/dify-agent /usr/local/bin/dify-agent RUN useradd --create-home --shell /bin/sh dify \ + && mkdir -p /workspace \ && chown dify:dify /home \ - && chown -R dify:dify /home/dify + && chown -R dify:dify /home/dify /workspace USER dify WORKDIR /home/dify diff --git a/dify-agent-runtime/internal/agentcli/httpclient.go b/dify-agent-runtime/internal/agentcli/httpclient.go index 93cd0d073e2..b22cf81359b 100644 --- a/dify-agent-runtime/internal/agentcli/httpclient.go +++ b/dify-agent-runtime/internal/agentcli/httpclient.go @@ -25,7 +25,10 @@ type HTTPClient struct { var errUploadRequestAborted = errors.New("upload request aborted") -const agentStubAuthorizationExpiredCode = "agent_stub_authorization_expired" +const ( + agentStubAuthorizationExpiredCode = "agent_stub_authorization_expired" + defaultUploadRequestTimeout = 180 * time.Second +) type agentStubHTTPError struct { statusCode int @@ -70,7 +73,7 @@ func openUploadSource(path string) (io.ReadCloser, error) { } func doUploadRequest(req *http.Request) (*http.Response, error) { - return (&http.Client{Timeout: 120 * time.Second}).Do(req) + return (&http.Client{Timeout: defaultUploadRequestTimeout}).Do(req) } // postJSON sends a POST request with JSON body and returns the response body. diff --git a/dify-agent-runtime/internal/agentcli/httpclient_test.go b/dify-agent-runtime/internal/agentcli/httpclient_test.go index 33d558f5378..c4ddf6c9830 100644 --- a/dify-agent-runtime/internal/agentcli/httpclient_test.go +++ b/dify-agent-runtime/internal/agentcli/httpclient_test.go @@ -20,6 +20,12 @@ import ( const fifoTestDeadline = 3 * time.Second +func TestDefaultUploadRequestTimeout(t *testing.T) { + if defaultUploadRequestTimeout != 180*time.Second { + t.Fatalf("defaultUploadRequestTimeout = %s, want 180s", defaultUploadRequestTimeout) + } +} + func receiveWithin[T any](ch <-chan T, timeout time.Duration) (T, bool) { timer := time.NewTimer(timeout) defer timer.Stop() diff --git a/dify-agent-runtime/internal/jobmode/mode.go b/dify-agent-runtime/internal/jobmode/mode.go new file mode 100644 index 00000000000..39741135b07 --- /dev/null +++ b/dify-agent-runtime/internal/jobmode/mode.go @@ -0,0 +1,29 @@ +// Package jobmode defines the execution modes shared by the shellctl server +// and runner. +package jobmode + +import "fmt" + +// Mode selects how the runner connects a job's standard streams. +type Mode string + +const ( + PTY Mode = "pty" + Stdio Mode = "stdio" +) + +// Parse validates a mode received at a process or API boundary. An empty value +// preserves the historical PTY behavior for callers that omit the mode. +func Parse(raw string) (Mode, error) { + if raw == "" { + return PTY, nil + } + + mode := Mode(raw) + switch mode { + case PTY, Stdio: + return mode, nil + default: + return "", fmt.Errorf("invalid job mode %q", raw) + } +} diff --git a/dify-agent-runtime/internal/jobmode/mode_test.go b/dify-agent-runtime/internal/jobmode/mode_test.go new file mode 100644 index 00000000000..b7ee65ffa0a --- /dev/null +++ b/dify-agent-runtime/internal/jobmode/mode_test.go @@ -0,0 +1,29 @@ +package jobmode + +import "testing" + +func TestParse(t *testing.T) { + tests := []struct { + name string + raw string + want Mode + wantErr bool + }{ + {name: "omitted", raw: "", want: PTY}, + {name: "pty", raw: "pty", want: PTY}, + {name: "stdio", raw: "stdio", want: Stdio}, + {name: "unknown", raw: "stdout", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := Parse(tt.raw) + if (err != nil) != tt.wantErr { + t.Fatalf("Parse(%q) error = %v, wantErr %v", tt.raw, err, tt.wantErr) + } + if got != tt.want { + t.Errorf("Parse(%q) = %q, want %q", tt.raw, got, tt.want) + } + }) + } +} diff --git a/dify-agent-runtime/internal/server/api.go b/dify-agent-runtime/internal/server/api.go index 67c4ee98772..0a6b34a9820 100644 --- a/dify-agent-runtime/internal/server/api.go +++ b/dify-agent-runtime/internal/server/api.go @@ -9,6 +9,8 @@ import ( "strconv" "strings" "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" ) // Handler creates the HTTP handler (mux) for the shellctl API. @@ -45,6 +47,12 @@ func handleRunJob(svc *Service) http.HandlerFunc { writeError(w, 400, "invalid_request", "script is required") return } + mode, err := jobmode.Parse(string(req.Mode)) + if err != nil { + writeError(w, 422, "validation_error", err.Error()) + return + } + req.Mode = mode // Validate env if req.Env != nil { for name, value := range req.Env { diff --git a/dify-agent-runtime/internal/server/api_test.go b/dify-agent-runtime/internal/server/api_test.go index b805143f26a..47d0977b4dd 100644 --- a/dify-agent-runtime/internal/server/api_test.go +++ b/dify-agent-runtime/internal/server/api_test.go @@ -152,6 +152,25 @@ func TestHealthzHandler(t *testing.T) { } } +func TestRunJobRejectsInvalidModeBeforeCallingService(t *testing.T) { + handler := handleRunJob(nil) + req := httptest.NewRequest("POST", "/v1/jobs/run", strings.NewReader(`{"script":"true","mode":"stdout"}`)) + w := httptest.NewRecorder() + + handler(w, req) + + if w.Code != http.StatusUnprocessableEntity { + t.Fatalf("expected 422, got %d", w.Code) + } + var result ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&result); err != nil { + t.Fatal(err) + } + if result.Error.Code != "validation_error" { + t.Errorf("expected validation_error, got %q", result.Error.Code) + } +} + func TestServerErrorFormat(t *testing.T) { err := NewServerError(422, "validation_error", "bad input") expected := "[422] validation_error: bad input" diff --git a/dify-agent-runtime/internal/server/db.go b/dify-agent-runtime/internal/server/db.go index a7eb6339b52..e6353e77e97 100644 --- a/dify-agent-runtime/internal/server/db.go +++ b/dify-agent-runtime/internal/server/db.go @@ -5,9 +5,19 @@ import ( "fmt" "time" + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" + _ "modernc.org/sqlite" ) +const latestSchemaVersion = 1 + +type schemaMigration func(*sql.Tx) error + +var schemaMigrations = []schemaMigration{ + migrateSchemaV1, +} + // JobStatusName represents the lifecycle states of a shellctl job. type JobStatusName string @@ -35,6 +45,7 @@ type JobRow struct { JobID string ScriptPath string OutputPath string + Mode jobmode.Mode Cwd string TerminalCols int TerminalRows int @@ -72,9 +83,63 @@ func (d *DB) Close() error { return d.db.Close() } -// InitSchema creates the jobs table if it does not exist. +// InitSchema creates the v0 baseline and applies pending schema migrations. func (d *DB) InitSchema() error { - _, err := d.db.Exec(` + var currentVersion int + if err := d.db.QueryRow("PRAGMA user_version").Scan(¤tVersion); err != nil { + return fmt.Errorf("read schema version: %w", err) + } + if currentVersion > latestSchemaVersion { + return fmt.Errorf( + "database schema version %d is newer than supported version %d", + currentVersion, + latestSchemaVersion, + ) + } + + if currentVersion == 0 { + if err := d.createSchemaV0(); err != nil { + return err + } + } + + for currentVersion < latestSchemaVersion { + targetVersion := currentVersion + 1 + if err := d.applySchemaMigration(targetVersion, schemaMigrations[targetVersion-1]); err != nil { + return err + } + currentVersion = targetVersion + } + return nil +} + +// applySchemaMigration commits migration SQL and PRAGMA user_version together, +// rolling back the entire version if either operation fails. +func (d *DB) applySchemaMigration(targetVersion int, migration schemaMigration) error { + tx, err := d.db.Begin() + if err != nil { + return fmt.Errorf("begin schema migration v%d: %w", targetVersion, err) + } + if err := migration(tx); err != nil { + _ = tx.Rollback() + return fmt.Errorf("apply schema migration v%d: %w", targetVersion, err) + } + if _, err := tx.Exec(fmt.Sprintf("PRAGMA user_version = %d", targetVersion)); err != nil { + _ = tx.Rollback() + return fmt.Errorf("record schema migration v%d: %w", targetVersion, err) + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit schema migration v%d: %w", targetVersion, err) + } + return nil +} + +func (d *DB) createSchemaV0() error { + tx, err := d.db.Begin() + if err != nil { + return fmt.Errorf("begin schema v0: %w", err) + } + if _, err := tx.Exec(` CREATE TABLE IF NOT EXISTS jobs ( job_id TEXT PRIMARY KEY, script_path TEXT NOT NULL, @@ -93,18 +158,29 @@ func (d *DB) InitSchema() error { ended_at TEXT, updated_at TEXT NOT NULL ) - `) + `); err != nil { + _ = tx.Rollback() + return fmt.Errorf("create schema v0: %w", err) + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit schema v0: %w", err) + } + return nil +} + +func migrateSchemaV1(tx *sql.Tx) error { + _, err := tx.Exec(`ALTER TABLE jobs ADD COLUMN mode TEXT NOT NULL DEFAULT 'pty'`) return err } // InsertJob inserts a new job row. Returns false if the job_id already exists. func (d *DB) InsertJob(row *JobRow) (bool, error) { _, err := d.db.Exec(` - INSERT INTO jobs (job_id, script_path, output_path, cwd, terminal_cols, terminal_rows, + INSERT INTO jobs (job_id, script_path, output_path, mode, cwd, terminal_cols, terminal_rows, status, session_name, pane_target, exit_code, reason, message, created_at, started_at, ended_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - row.JobID, row.ScriptPath, row.OutputPath, row.Cwd, + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + row.JobID, row.ScriptPath, row.OutputPath, string(row.Mode), row.Cwd, row.TerminalCols, row.TerminalRows, string(row.Status), row.SessionName, row.PaneTarget, row.ExitCode, row.Reason, row.Message, @@ -123,7 +199,7 @@ func (d *DB) InsertJob(row *JobRow) (bool, error) { // GetJob retrieves a single job row by ID. func (d *DB) GetJob(jobID string) (*JobRow, error) { row := d.db.QueryRow(` - SELECT job_id, script_path, output_path, cwd, terminal_cols, terminal_rows, + SELECT job_id, script_path, output_path, mode, cwd, terminal_cols, terminal_rows, status, session_name, pane_target, exit_code, reason, message, created_at, started_at, ended_at, updated_at FROM jobs WHERE job_id = ?`, jobID) @@ -136,7 +212,7 @@ func (d *DB) ListJobs(statuses []JobStatusName) ([]*JobRow, error) { var err error if len(statuses) == 0 { - rows, err = d.db.Query(`SELECT job_id, script_path, output_path, cwd, + rows, err = d.db.Query(`SELECT job_id, script_path, output_path, mode, cwd, terminal_cols, terminal_rows, status, session_name, pane_target, exit_code, reason, message, created_at, started_at, ended_at, updated_at FROM jobs ORDER BY created_at DESC`) @@ -150,7 +226,7 @@ func (d *DB) ListJobs(statuses []JobStatusName) ([]*JobRow, error) { } placeholders += "?" } - query := fmt.Sprintf(`SELECT job_id, script_path, output_path, cwd, + query := fmt.Sprintf(`SELECT job_id, script_path, output_path, mode, cwd, terminal_cols, terminal_rows, status, session_name, pane_target, exit_code, reason, message, created_at, started_at, ended_at, updated_at FROM jobs WHERE status IN (%s) ORDER BY created_at DESC`, placeholders) @@ -313,9 +389,9 @@ type TransitionOpts struct { func scanJobRow(row *sql.Row) (*JobRow, error) { var jr JobRow - var status string + var mode, status string err := row.Scan( - &jr.JobID, &jr.ScriptPath, &jr.OutputPath, &jr.Cwd, + &jr.JobID, &jr.ScriptPath, &jr.OutputPath, &mode, &jr.Cwd, &jr.TerminalCols, &jr.TerminalRows, &status, &jr.SessionName, &jr.PaneTarget, &jr.ExitCode, &jr.Reason, &jr.Message, @@ -327,15 +403,16 @@ func scanJobRow(row *sql.Row) (*JobRow, error) { if err != nil { return nil, err } + jr.Mode = jobmode.Mode(mode) jr.Status = JobStatusName(status) return &jr, nil } func scanJobRows(rows *sql.Rows) (*JobRow, error) { var jr JobRow - var status string + var mode, status string err := rows.Scan( - &jr.JobID, &jr.ScriptPath, &jr.OutputPath, &jr.Cwd, + &jr.JobID, &jr.ScriptPath, &jr.OutputPath, &mode, &jr.Cwd, &jr.TerminalCols, &jr.TerminalRows, &status, &jr.SessionName, &jr.PaneTarget, &jr.ExitCode, &jr.Reason, &jr.Message, @@ -344,6 +421,7 @@ func scanJobRows(rows *sql.Rows) (*JobRow, error) { if err != nil { return nil, err } + jr.Mode = jobmode.Mode(mode) jr.Status = JobStatusName(status) return &jr, nil } diff --git a/dify-agent-runtime/internal/server/db_test.go b/dify-agent-runtime/internal/server/db_test.go index c3bbc2c3027..2f1e236f4ac 100644 --- a/dify-agent-runtime/internal/server/db_test.go +++ b/dify-agent-runtime/internal/server/db_test.go @@ -1,9 +1,13 @@ package server import ( + "database/sql" + "errors" "os" "path/filepath" "testing" + + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" ) func TestJobStatusIsTerminal(t *testing.T) { @@ -40,6 +44,7 @@ func TestOpenDBAndInitSchema(t *testing.T) { JobID: "test-job-1", ScriptPath: "jobs/test-job-1/script", OutputPath: "jobs/test-job-1/output.log", + Mode: jobmode.PTY, Cwd: "/tmp", TerminalCols: 80, TerminalRows: 24, @@ -66,6 +71,104 @@ func TestOpenDBAndInitSchema(t *testing.T) { if ok { t.Error("expected duplicate insert to return ok=false") } + if got := schemaVersion(t, db); got != latestSchemaVersion { + t.Errorf("schema version = %d, want %d", got, latestSchemaVersion) + } +} + +func TestInitSchemaMigratesV0JobsToPTY(t *testing.T) { + db := openTestDB(t, t.TempDir()) + defer func() { _ = db.Close() }() + + if err := db.createSchemaV0(); err != nil { + t.Fatalf("createSchemaV0: %v", err) + } + _, err := db.db.Exec(` + INSERT INTO jobs ( + job_id, script_path, output_path, cwd, terminal_cols, terminal_rows, + status, session_name, pane_target, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + "legacy-job", "jobs/legacy-job/script", "jobs/legacy-job/output.log", "/tmp", + 80, 24, "created", "shellctl-legacy-job", "shellctl-legacy-job:0.0", + "2025-01-01T00:00:00Z", "2025-01-01T00:00:00Z", + ) + if err != nil { + t.Fatalf("insert legacy job: %v", err) + } + + if err := db.InitSchema(); err != nil { + t.Fatalf("InitSchema: %v", err) + } + row, err := db.GetJob("legacy-job") + if err != nil { + t.Fatalf("GetJob: %v", err) + } + if row.Mode != jobmode.PTY { + t.Errorf("legacy job mode = %q, want %q", row.Mode, jobmode.PTY) + } + if got := schemaVersion(t, db); got != latestSchemaVersion { + t.Errorf("schema version = %d, want %d", got, latestSchemaVersion) + } +} + +func TestInitSchemaAtLatestVersionIsIdempotent(t *testing.T) { + db := setupTestDB(t, t.TempDir()) + defer func() { _ = db.Close() }() + + if err := db.InitSchema(); err != nil { + t.Fatalf("second InitSchema: %v", err) + } + if got := schemaVersion(t, db); got != latestSchemaVersion { + t.Errorf("schema version = %d, want %d", got, latestSchemaVersion) + } +} + +func TestInitSchemaRejectsNewerDatabaseWithoutDDL(t *testing.T) { + db := openTestDB(t, t.TempDir()) + defer func() { _ = db.Close() }() + if _, err := db.db.Exec("PRAGMA user_version = 2"); err != nil { + t.Fatalf("set future schema version: %v", err) + } + + if err := db.InitSchema(); err == nil { + t.Fatal("InitSchema unexpectedly accepted a newer schema") + } + var tableCount int + if err := db.db.QueryRow(`SELECT count(*) FROM sqlite_master WHERE type='table' AND name='jobs'`).Scan(&tableCount); err != nil { + t.Fatalf("query jobs table: %v", err) + } + if tableCount != 0 { + t.Errorf("jobs table count = %d, want 0", tableCount) + } +} + +func TestApplySchemaMigrationFailureRollsBackDDLAndVersion(t *testing.T) { + db := openTestDB(t, t.TempDir()) + defer func() { _ = db.Close() }() + if err := db.createSchemaV0(); err != nil { + t.Fatalf("createSchemaV0: %v", err) + } + + sentinel := errors.New("sentinel migration failure") + err := db.applySchemaMigration(1, func(tx *sql.Tx) error { + if _, err := tx.Exec(`ALTER TABLE jobs ADD COLUMN rollback_probe TEXT`); err != nil { + return err + } + return sentinel + }) + if !errors.Is(err, sentinel) { + t.Fatalf("migration error = %v, want sentinel failure", err) + } + var probeColumns int + if err := db.db.QueryRow(`SELECT count(*) FROM pragma_table_info('jobs') WHERE name = 'rollback_probe'`).Scan(&probeColumns); err != nil { + t.Fatalf("query rollback probe column: %v", err) + } + if probeColumns != 0 { + t.Errorf("rollback_probe column count = %d, want 0", probeColumns) + } + if got := schemaVersion(t, db); got != 0 { + t.Errorf("schema version = %d, want 0", got) + } } func TestGetJob(t *testing.T) { @@ -88,6 +191,9 @@ func TestGetJob(t *testing.T) { if row.TerminalCols != 80 { t.Errorf("expected cols=80, got %d", row.TerminalCols) } + if row.Mode != jobmode.PTY { + t.Errorf("expected mode=pty, got %s", row.Mode) + } } func TestGetJobNotFound(t *testing.T) { @@ -243,6 +349,7 @@ func TestRecordRunnerExitIdempotent(t *testing.T) { exitCode := 10 row := &JobRow{ JobID: "job-exit-2", ScriptPath: "x", OutputPath: "y", Cwd: "/tmp", + Mode: jobmode.PTY, TerminalCols: 80, TerminalRows: 24, Status: StatusExited, SessionName: "s", PaneTarget: "p", ExitCode: &exitCode, CreatedAt: "2025-01-01T00:00:00Z", UpdatedAt: "2025-01-01T00:00:00Z", @@ -267,24 +374,40 @@ func TestRecordRunnerExitIdempotent(t *testing.T) { // Helpers func setupTestDB(t *testing.T, dir string) *DB { + t.Helper() + db := openTestDB(t, dir) + if err := db.InitSchema(); err != nil { + t.Fatalf("InitSchema: %v", err) + } + return db +} + +func openTestDB(t *testing.T, dir string) *DB { t.Helper() dbPath := filepath.Join(dir, "shellctl.db") db, err := OpenDB(dbPath, 5000) if err != nil { t.Fatalf("OpenDB: %v", err) } - if err := db.InitSchema(); err != nil { - t.Fatalf("InitSchema: %v", err) - } return db } +func schemaVersion(t *testing.T, db *DB) int { + t.Helper() + var version int + if err := db.db.QueryRow("PRAGMA user_version").Scan(&version); err != nil { + t.Fatalf("read schema version: %v", err) + } + return version +} + func insertTestJob(t *testing.T, db *DB, jobID string, status JobStatusName) { t.Helper() row := &JobRow{ JobID: jobID, ScriptPath: "jobs/" + jobID + "/script", OutputPath: "jobs/" + jobID + "/output.log", + Mode: jobmode.PTY, Cwd: "/tmp", TerminalCols: 80, TerminalRows: 24, diff --git a/dify-agent-runtime/internal/server/service.go b/dify-agent-runtime/internal/server/service.go index a859aa3474d..693d3c16fcc 100644 --- a/dify-agent-runtime/internal/server/service.go +++ b/dify-agent-runtime/internal/server/service.go @@ -11,6 +11,8 @@ import ( "strings" "sync" "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" ) // Service is the core job lifecycle manager backed by SQLite and tmux. @@ -105,7 +107,9 @@ func (s *Service) StartBackgroundGC() { }() } -// StartBackgroundPipeMonitor starts the periodic pipe health check goroutine. +// StartBackgroundPipeMonitor periodically reconciles mode-aware runtime state: +// PTY pane pipe/drain state, and stdio session state plus completion +// materialization from exit artifacts. func (s *Service) StartBackgroundPipeMonitor() { ctx, cancel := context.WithCancel(context.Background()) s.cancelMon = cancel @@ -194,6 +198,7 @@ func (s *Service) RunJob(req *RunJobRequest) (*JobResult, error) { JobID: jobID, ScriptPath: fmt.Sprintf("jobs/%s/script", jobID), OutputPath: fmt.Sprintf("jobs/%s/output.log", jobID), + Mode: req.Mode, Cwd: cwd, TerminalCols: cols, TerminalRows: rows, @@ -220,9 +225,10 @@ func (s *Service) RunJob(req *RunJobRequest) (*JobResult, error) { Target: StatusStarting, }) - // Create tmux session and enable output pipe + // Start the mode-specific runtime: PTY manages pane piping and drain, while + // stdio tracks the session and materializes completion from exit artifacts. log.Printf("RunJob [%s]: starting job, cwd=%s", jobID, cwd) - startErr := s.startJob(jobID, jobDir, cwd, cols, rows) + startErr := s.startJob(jobID, jobDir, cwd, cols, rows, req.Mode) if startErr != nil { log.Printf("RunJob [%s]: start failed: %v", jobID, startErr) reason := "start_failed" @@ -251,24 +257,26 @@ func (s *Service) RunJob(req *RunJobRequest) (*JobResult, error) { }) } -func (s *Service) startJob(jobID, jobDir, cwd string, cols, rows int) error { +func (s *Service) startJob(jobID, jobDir, cwd string, cols, rows int, mode jobmode.Mode) error { log.Printf("startJob [%s]: creating tmux session", jobID) - if err := s.tmux.CreateJobSession(jobID, jobDir, cwd, cols, rows); err != nil { + if err := s.tmux.CreateJobSession(jobID, jobDir, cwd, cols, rows, mode); err != nil { log.Printf("startJob [%s]: tmux session failed: %v", jobID, err) return err } pipeReadyPath := filepath.Join(jobDir, ".pipe-ready") - log.Printf("startJob [%s]: enabling output pipe", jobID) - if err := s.tmux.EnableOutputPipe(jobID, jobDir, pipeReadyPath); err != nil { - log.Printf("startJob [%s]: pipe-pane failed: %v", jobID, err) - return err - } + if mode == jobmode.PTY { + log.Printf("startJob [%s]: enabling output pipe", jobID) + if err := s.tmux.EnableOutputPipe(jobID, jobDir, pipeReadyPath); err != nil { + log.Printf("startJob [%s]: pipe-pane failed: %v", jobID, err) + return err + } - // Wait for pipe ready handshake - if err := s.waitForPipeReady(jobID, pipeReadyPath); err != nil { - log.Printf("startJob [%s]: pipe-ready timeout: %v", jobID, err) - return err + // Wait for pipe ready handshake + if err := s.waitForPipeReady(jobID, pipeReadyPath); err != nil { + log.Printf("startJob [%s]: pipe-ready timeout: %v", jobID, err) + return err + } } // Open start gate @@ -285,8 +293,9 @@ func (s *Service) startJob(jobID, jobDir, cwd string, cols, rows int) error { RequireExitCodeNull: true, }) - // Clean up ready file - _ = os.Remove(pipeReadyPath) + if mode == jobmode.PTY { + _ = os.Remove(pipeReadyPath) + } return nil } @@ -425,11 +434,15 @@ func (s *Service) TailJob(jobID string, outputLimit int) (*JobResult, error) { // GetJobStatus materializes the current status from SQLite + live tmux state. func (s *Service) GetJobStatus(jobID string) (*JobStatusView, error) { - sessionExists, pipeActive, err := s.liveRuntimeState(jobID) + row, err := s.db.GetJob(jobID) if err != nil { return nil, err } - return s.materializeStatusView(jobID, sessionExists, pipeActive) + sessionExists, pipeActive, err := s.liveRuntimeState(row) + if err != nil { + return nil, err + } + return s.materializeStatusView(row, sessionExists, pipeActive) } // ListJobs returns recent jobs, optionally filtered by status. @@ -474,6 +487,13 @@ func (s *Service) SendInput(jobID string, req *InputJobRequest) (*JobResult, err if view.Done { return nil, NewServerError(409, "job_not_running", fmt.Sprintf("Job %s is already terminal", jobID)) } + row, err := s.db.GetJob(jobID) + if err != nil { + return nil, err + } + if row.Mode == jobmode.Stdio { + return nil, NewServerError(409, "input_unsupported", "stdio jobs do not support input") + } if err := s.tmux.SendInput(jobID, req.Text); err != nil { // Check if job became terminal in the meantime @@ -594,7 +614,8 @@ func (s *Service) GCOnce() error { return nil } -// CheckRunningJobsPipeHealth fails running jobs whose pipe died. +// CheckRunningJobsPipeHealth reconciles PTY pane pipe/drain state and stdio +// session state, materializing completion from exit artifacts. func (s *Service) CheckRunningJobsPipeHealth() { rows, _ := s.db.ListJobs([]JobStatusName{StatusRunning}) for _, row := range rows { @@ -602,12 +623,8 @@ func (s *Service) CheckRunningJobsPipeHealth() { } } -func (s *Service) materializeStatusView(jobID string, sessionExists bool, pipeActive *bool) (*JobStatusView, error) { - row, err := s.db.GetJob(jobID) - if err != nil { - return nil, err - } - +func (s *Service) materializeStatusView(row *JobRow, sessionExists bool, pipeActive *bool) (*JobStatusView, error) { + jobID := row.JobID status := row.Status if status.IsTerminal() { @@ -631,14 +648,14 @@ func (s *Service) materializeStatusView(jobID string, sessionExists bool, pipeAc }); err == nil { row = r } - } else if exit := s.drainedNormalExitMetadata(jobID); exit != nil { - // Recover from drained exit artifacts + } else if exit := s.completedExitMetadata(row); exit != nil { + // Recover from mode-specific completed exit artifacts. _ = s.db.RecordRunnerExit(jobID, exit.exitCode, exit.endedAt) if r, err := s.db.GetJob(jobID); err == nil { row = r } } else if sessionExists { - if pipeActive != nil && !*pipeActive { + if row.Mode == jobmode.PTY && pipeActive != nil && !*pipeActive { s.mu.Lock() isStarting := s.startingJobs[jobID] s.mu.Unlock() @@ -670,7 +687,7 @@ func (s *Service) materializeStatusView(jobID string, sessionExists bool, pipeAc } } else { // No session - if s.normalExitCommitPending(jobID) { + if s.normalExitCommitPending(row) { // Wait for pipe drain finalizer } else { s.mu.Lock() @@ -694,15 +711,18 @@ func (s *Service) materializeStatusView(jobID string, sessionExists bool, pipeAc return s.statusViewFromRow(row), nil } -func (s *Service) liveRuntimeState(jobID string) (bool, *bool, error) { - exists, err := s.tmux.SessionExists(JobSessionName(jobID)) +func (s *Service) liveRuntimeState(row *JobRow) (bool, *bool, error) { + exists, err := s.tmux.SessionExists(row.SessionName) if err != nil { return false, nil, err } if !exists { return false, nil, nil } - active, err := s.tmux.IsOutputPipeActive(jobID) + if row.Mode == jobmode.Stdio { + return true, nil, nil + } + active, err := s.tmux.IsOutputPipeActive(row.JobID) if err != nil { return false, nil, err } @@ -717,13 +737,16 @@ type exitMetadata struct { endedAt string } -func (s *Service) drainedNormalExitMetadata(jobID string) *exitMetadata { - jobDir := filepath.Join(s.config.JobsDir(), jobID) +func (s *Service) completedExitMetadata(row *JobRow) *exitMetadata { + jobDir := filepath.Join(s.config.JobsDir(), row.JobID) drainedPath := filepath.Join(jobDir, ".pipe-drained") exitCodePath := filepath.Join(jobDir, "runner-exit-code") endedAtPath := filepath.Join(jobDir, "runner-ended-at") - if !fileExists(drainedPath) || !fileExists(exitCodePath) || !fileExists(endedAtPath) { + if row.Mode == jobmode.PTY && !fileExists(drainedPath) { + return nil + } + if !fileExists(exitCodePath) || !fileExists(endedAtPath) { return nil } @@ -750,8 +773,11 @@ func (s *Service) drainedNormalExitMetadata(jobID string) *exitMetadata { return &exitMetadata{exitCode: code, endedAt: endedAtStr} } -func (s *Service) normalExitCommitPending(jobID string) bool { - jobDir := filepath.Join(s.config.JobsDir(), jobID) +func (s *Service) normalExitCommitPending(row *JobRow) bool { + if row.Mode != jobmode.PTY { + return false + } + jobDir := filepath.Join(s.config.JobsDir(), row.JobID) return fileExists(filepath.Join(jobDir, "runner-exit-code")) && fileExists(filepath.Join(jobDir, "runner-ended-at")) && !fileExists(filepath.Join(jobDir, ".pipe-drained")) && diff --git a/dify-agent-runtime/internal/server/service_test.go b/dify-agent-runtime/internal/server/service_test.go new file mode 100644 index 00000000000..1f2405c86d4 --- /dev/null +++ b/dify-agent-runtime/internal/server/service_test.go @@ -0,0 +1,111 @@ +package server + +import ( + "os" + "path/filepath" + "testing" + + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" +) + +func TestMaterializeStdioStatusDoesNotRequirePanePipe(t *testing.T) { + service, row := setupModeTestService(t, jobmode.Stdio, StatusRunning) + pipeInactive := false + + view, err := service.materializeStatusView(row, true, &pipeInactive) + if err != nil { + t.Fatalf("materializeStatusView: %v", err) + } + if view.Status != StatusRunning { + t.Errorf("status = %q, want %q", view.Status, StatusRunning) + } +} + +func TestMaterializeStdioStatusUsesExitArtifactsWithoutPipeDrainMarker(t *testing.T) { + service, row := setupModeTestService(t, jobmode.Stdio, StatusRunning) + jobDir := filepath.Join(service.config.JobsDir(), row.JobID) + if err := os.WriteFile(filepath.Join(jobDir, "runner-exit-code"), []byte("7\n"), 0600); err != nil { + t.Fatal(err) + } + const endedAt = "2026-08-04T10:00:00Z" + if err := os.WriteFile(filepath.Join(jobDir, "runner-ended-at"), []byte(endedAt+"\n"), 0600); err != nil { + t.Fatal(err) + } + + view, err := service.materializeStatusView(row, false, nil) + if err != nil { + t.Fatalf("materializeStatusView: %v", err) + } + if view.Status != StatusExited || !view.Done { + t.Errorf("view = status %q done %v, want exited and done", view.Status, view.Done) + } + if view.ExitCode == nil || *view.ExitCode != 7 { + t.Errorf("exit code = %v, want 7", view.ExitCode) + } + if view.EndedAt == nil || *view.EndedAt != endedAt { + t.Errorf("ended_at = %v, want %s", view.EndedAt, endedAt) + } +} + +func TestMaterializeStdioStatusMarksMissingSessionWithIncompleteArtifactsLost(t *testing.T) { + service, row := setupModeTestService(t, jobmode.Stdio, StatusRunning) + jobDir := filepath.Join(service.config.JobsDir(), row.JobID) + if err := os.WriteFile(filepath.Join(jobDir, "runner-exit-code"), []byte("0\n"), 0600); err != nil { + t.Fatal(err) + } + + view, err := service.materializeStatusView(row, false, nil) + if err != nil { + t.Fatalf("materializeStatusView: %v", err) + } + if view.Status != StatusLost { + t.Errorf("status = %q, want %q", view.Status, StatusLost) + } +} + +func setupModeTestService(t *testing.T, mode jobmode.Mode, status JobStatusName) (*Service, *JobRow) { + t.Helper() + stateDir := t.TempDir() + config := DefaultConfig() + config.StateDir = stateDir + config.RuntimeDir = filepath.Join(stateDir, "runtime") + if err := os.MkdirAll(config.JobsDir(), 0700); err != nil { + t.Fatal(err) + } + db := setupTestDB(t, stateDir) + t.Cleanup(func() { _ = db.Close() }) + service := NewService(config) + service.db = db + + const jobID = "mode-test-job" + jobDir := filepath.Join(config.JobsDir(), jobID) + if err := os.MkdirAll(jobDir, 0700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(jobDir, "output.log"), nil, 0600); err != nil { + t.Fatal(err) + } + row := &JobRow{ + JobID: jobID, + ScriptPath: "jobs/mode-test-job/script", + OutputPath: "jobs/mode-test-job/output.log", + Mode: mode, + Cwd: "/tmp", + TerminalCols: 80, + TerminalRows: 24, + Status: status, + SessionName: JobSessionName(jobID), + PaneTarget: JobPaneTarget(jobID), + CreatedAt: "2026-08-04T09:00:00Z", + UpdatedAt: "2026-08-04T09:00:00Z", + } + inserted, err := db.InsertJob(row) + if err != nil || !inserted { + t.Fatalf("InsertJob: inserted=%v err=%v", inserted, err) + } + persisted, err := db.GetJob(jobID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + return service, persisted +} diff --git a/dify-agent-runtime/internal/server/tmux.go b/dify-agent-runtime/internal/server/tmux.go index dc5cfac9565..839b06b229e 100644 --- a/dify-agent-runtime/internal/server/tmux.go +++ b/dify-agent-runtime/internal/server/tmux.go @@ -7,6 +7,8 @@ import ( "os/exec" "path/filepath" "strings" + + "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" ) // TmuxController manages tmux sessions for shellctl jobs via a dedicated socket. @@ -82,9 +84,13 @@ func (t *TmuxController) IsOutputPipeActive(jobID string) (*bool, error) { } // CreateJobSession creates a new tmux session for a job. -func (t *TmuxController) CreateJobSession(jobID, jobDir, cwd string, cols, rows int) error { +func (t *TmuxController) CreateJobSession( + jobID, jobDir, cwd string, + cols, rows int, + mode jobmode.Mode, +) error { runnerCmd := shellJoin([]string{ - t.config.RunnerPath(), jobDir, jobID, cwd, + t.config.RunnerPath(), jobDir, jobID, cwd, string(mode), }) result, err := t.runTmuxNoCheck( "-f", "/dev/null", diff --git a/dify-agent-runtime/internal/server/types.go b/dify-agent-runtime/internal/server/types.go index 28eecc37b86..b23c9c1e3f8 100644 --- a/dify-agent-runtime/internal/server/types.go +++ b/dify-agent-runtime/internal/server/types.go @@ -1,11 +1,14 @@ package server +import "github.com/langgenius/dify/dify-agent-runtime/internal/jobmode" + // RunJobRequest is the HTTP request body for POST /v1/jobs/run. type RunJobRequest struct { Script string `json:"script"` Cwd *string `json:"cwd,omitempty"` Env map[string]string `json:"env,omitempty"` Terminal *TerminalSize `json:"terminal,omitempty"` + Mode jobmode.Mode `json:"mode,omitempty"` Timeout float64 `json:"timeout,omitempty"` OutputLimit int `json:"output_limit,omitempty"` IdleFlushSeconds float64 `json:"idle_flush_seconds,omitempty"` diff --git a/dify-agent-runtime/tests/acceptance_test.go b/dify-agent-runtime/tests/acceptance_test.go index aed05e6ac9d..192d75729fa 100644 --- a/dify-agent-runtime/tests/acceptance_test.go +++ b/dify-agent-runtime/tests/acceptance_test.go @@ -170,6 +170,84 @@ func TestRunSimpleScript(t *testing.T) { } } +func TestPTYModesMergeStdoutAndStderr(t *testing.T) { + for _, tgt := range targets() { + for _, tc := range []struct { + name string + mode string + }{ + {name: "default"}, + {name: "explicit", mode: "pty"}, + } { + t.Run(tgt.name+"/"+tc.name, func(t *testing.T) { + payload := map[string]any{ + "script": "printf 'stdout-marker\\n'; printf 'stderr-marker\\n' >&2", + "timeout": 10, + } + if tc.mode != "" { + payload["mode"] = tc.mode + } + result := runJob(t, tgt, payload) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "stdout-marker") || !strings.Contains(output, "stderr-marker") { + t.Fatalf("PTY output did not merge stdout and stderr: %q", output) + } + }) + } + } +} + +func TestRunStdioSeparatesStdoutAndStderr(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "printf '{\"ok\":true}'\nprintf 'warning' >&2", + "mode": "stdio", + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + if output := result["output"].(string); output != `{"ok":true}` { + t.Errorf("stdio output = %q, want stdout-only JSON", output) + } + }) + } +} + +func TestStdioInputIsRejected(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "sleep 60", + "mode": "stdio", + "timeout": 0.1, + }) + jobID := result["job_id"].(string) + defer func() { + resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/terminate", jobID), map[string]any{"grace_seconds": 0}, true) + resp.Body.Close() + }() + + resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/input", jobID), map[string]any{ + "text": "ignored\n", + "offset": 0, + "timeout": 1, + }, true) + assertStatus(t, resp, http.StatusConflict) + body := readBody(t, resp) + var failure map[string]map[string]string + if err := json.Unmarshal(body, &failure); err != nil { + t.Fatal(err) + } + if code := failure["error"]["code"]; code != "input_unsupported" { + t.Errorf("error code = %q, want input_unsupported", code) + } + }) + } +} + func TestRunWithEnv(t *testing.T) { for _, tgt := range targets() { t.Run(tgt.name, func(t *testing.T) { @@ -431,10 +509,12 @@ func TestSendInput(t *testing.T) { "timeout": 2, // Will timeout waiting for input }) jobID := result["job_id"].(string) + t.Cleanup(func() { + cleanupJobBestEffort(tgt, jobID) + }) if result["done"] == true { - // Already finished (possible race), skip - t.Skip("job completed before input could be sent") + t.Fatalf("interactive PTY job completed before input was sent: %#v", result) } // Send input @@ -451,7 +531,7 @@ func TestSendInput(t *testing.T) { json.Unmarshal(body, &inputResult) output := inputResult["output"].(string) if !strings.Contains(output, "got:hello-input") { - t.Logf("output after input: %q (may need more wait time)", output) + t.Fatalf("input result did not contain command echo: %q", output) } }) } @@ -551,19 +631,27 @@ func TestLandlockCanReadSystemBinaries(t *testing.T) { } } -func TestLandlockCanWriteTmpdir(t *testing.T) { +func TestLandlockUsesWorkspaceAsTempSpace(t *testing.T) { for _, tgt := range targets() { t.Run(tgt.name, func(t *testing.T) { - // TMPDIR ($CWD/.tmp) should be writable; /tmp should be denied. + // The workspace itself is cwd and temp space; shared /tmp remains denied. result := runJob(t, tgt, map[string]any{ - "script": "echo TMPDIR=$TMPDIR && touch $TMPDIR/landlock-tmp-test && echo tmpdir_ok && touch /tmp/landlock-denied 2>&1; echo tmp_exit=$?", - "env": map[string]string{"HOME": "/home/dify"}, + "script": "test \"$PWD\" = /workspace && test \"$TMPDIR\" = /workspace && test \"$TMP\" = /workspace && test \"$TEMP\" = /workspace && " + + "touch \"$TMPDIR/landlock-tmp-test\" && echo workspace_temp_ok; " + + "touch /tmp/landlock-denied 2>&1; echo tmp_exit=$?", + "cwd": "/workspace", + "env": map[string]string{ + "HOME": "/home/dify", + "TMPDIR": "/tmp", + "TMP": "/tmp", + "TEMP": "/tmp", + }, "timeout": 10, }) assertJobDone(t, result) output := result["output"].(string) - if !strings.Contains(output, "tmpdir_ok") { - t.Errorf("expected write to $TMPDIR to succeed, got %q", output) + if !strings.Contains(output, "workspace_temp_ok") { + t.Errorf("expected workspace temp checks to pass, got %q", output) } if !strings.Contains(output, "tmp_exit=1") && !strings.Contains(output, "Permission denied") { t.Errorf("expected write to /tmp to be denied, got %q", output) @@ -572,6 +660,48 @@ func TestLandlockCanWriteTmpdir(t *testing.T) { } } +func TestRunnerDoesNotCreateCwdTmpDirectory(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + freshCwd := fmt.Sprintf("/workspace/no-auto-tmp-%s-%d", tgt.name, time.Now().UnixNano()) + setup := runJob(t, tgt, map[string]any{ + "script": "mkdir -p -- \"$FRESH_CWD\"", + "cwd": "/workspace", + "env": map[string]string{ + "HOME": "/home/dify", + "FRESH_CWD": freshCwd, + }, + "timeout": 10, + }) + assertJobDone(t, setup) + assertExitCode(t, setup, 0) + + t.Cleanup(func() { + cleanup := runJob(t, tgt, map[string]any{ + "script": "rm -rf -- \"$FRESH_CWD\"", + "cwd": "/workspace", + "env": map[string]string{ + "HOME": "/home/dify", + "FRESH_CWD": freshCwd, + }, + "timeout": 10, + }) + assertJobDone(t, cleanup) + assertExitCode(t, cleanup, 0) + }) + + result := runJob(t, tgt, map[string]any{ + "script": "test ! -e \"$PWD/.tmp\"", + "cwd": freshCwd, + "env": map[string]string{"HOME": "/home/dify"}, + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + }) + } +} + func TestLandlockCannotWriteOutsideHome(t *testing.T) { for _, tgt := range targets() { t.Run(tgt.name, func(t *testing.T) { @@ -720,6 +850,25 @@ func doPost(t *testing.T, tgt target, path string, payload map[string]any, withA return resp } +func cleanupJobBestEffort(tgt target, jobID string) { + client := &http.Client{Timeout: 5 * time.Second} + body, _ := json.Marshal(map[string]any{"grace_seconds": 0}) + req, _ := http.NewRequest("POST", fmt.Sprintf("%s/v1/jobs/%s/terminate", tgt.baseURL, jobID), bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+authToken) + if resp, err := client.Do(req); err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } + + req, _ = http.NewRequest("DELETE", fmt.Sprintf("%s/v1/jobs/%s?force=true&grace_seconds=0", tgt.baseURL, jobID), nil) + req.Header.Set("Authorization", "Bearer "+authToken) + if resp, err := client.Do(req); err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } +} + func readBody(t *testing.T, resp *http.Response) []byte { t.Helper() defer resp.Body.Close() diff --git a/dify-agent/.example.env b/dify-agent/.example.env index d89310a9e41..5aeb0032157 100644 --- a/dify-agent/.example.env +++ b/dify-agent/.example.env @@ -60,6 +60,8 @@ DIFY_AGENT_STUB_API_BASE_URL=http://localhost:5050/agent-stub DIFY_AGENT_SANDBOX_FILES_BASE_URL=http://localhost:5001 # Maximum Agent Stub upload size in MiB; forwarded to Dify API as a signed byte limit. DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT=50 +# Shell command deadline for converting a Binding file to a ToolFile. +DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS=210 # Server-wide root secret used to derive Agent Stub JWE keys. # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. diff --git a/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md b/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md index c3bae6fb3c6..c675898bd60 100644 --- a/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md +++ b/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md @@ -118,10 +118,12 @@ product use without performing network I/O inside the caller's transaction. Product lifecycle paths commit this transition synchronously. After the transaction commits, one Celery task asks Dify Agent to destroy the physical resources. A successful collector deletes the corresponding ledger row. If a -collector raises, the task logs the tenant, resource type, and resource ID, -re-raises the exception so collection stops and Celery records the task as -failed, and leaves the RETIRED row intact. No automatic retry or reconciliation -is performed. +collector raises, the task logs the tenant, resource type, and resource ID and +continues with the other independent resources in the batch. After all resources +have been attempted, any failure makes the Celery task fail and prevents Agent +aggregate deletion. Failed RETIRED rows remain available for a later retry. A +failure to publish the Celery task is also propagated to the product caller. No +automatic retry or reconciliation is performed. The unified `collect_agent_resources` task is registered on normal Celery workers and explicitly uses the existing `retention` queue. Standard workers @@ -130,14 +132,26 @@ is required. At a Workflow terminal event, the graph layer synchronously retires and commits the run's Workspaces before enqueueing collection. When a Workflow change may orphan Workflow-only Agents, the main product transaction commits first; a fresh session then rechecks effective ownership and retires only Agents -that remain unowned. +that remain unowned. An effective reference is a binding in a normal App's +current draft or current published Workflow. This ownership check applies only +to implicit retirement of Workflow-only Agents. Explicit deletion of a roster +Agent or Agent App proceeds even while Workflows reference it. Retiring a final Binding also retires its Workspace. Workspace collection destroys the physical Workspace through one Binding and then collects remaining materialized Homes. Home Snapshots are retired when their owning Agent is -retired and are collected only after no draft or config snapshot references -them. Celery performs physical collection only; it does not decide or perform -the initial retirement. Dify Agent itself remains stateless. +retired. `RETIRED` is the sole physical-deletion condition for a Home Snapshot; +Draft and Config Snapshot references are historical pointers and do not keep it +alive. After every external resource in a deletion batch succeeds, Dify API +hard-deletes the archived Agent together with its Drafts, Config Snapshots, +Config Revisions, debug-conversation mappings, and resource ledgers in one +database transaction. Workflow Agent bindings belong to +their Workflows and remain unchanged, so they may hold a dangling Agent ID after +explicit deletion. Dify Agent itself remains stateless. + +A `RETIRED` Workspace without a `RETIRED` Binding cannot identify a backend +participant through which to destroy the Workspace. That state is a lifecycle +invariant violation and fails collection instead of being logged as success. There is currently no age-based TTL, periodic GC, or global orphan reconciler. Backend destroy operations are idempotent where supported. Dify API does not @@ -173,12 +187,20 @@ directly from the runtime to Dify's existing ToolFile endpoint. Dify Agent returns only the canonical ToolFile reference and releases the lease before Dify API signs a browser URL. +The default Binding-download deadline chain leaves each caller time to receive +and normalize the lower layer's result: the sandbox CLI upload is 180 seconds, +`DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS` is 210 seconds, +Dify API's `AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS` is 240 seconds, +and `DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` is 3600 seconds. + `RuntimeLayout.home_dir` and `RuntimeLayout.workspace_dir` are canonical paths inside the backend execution namespace. They are not host paths, product ids, or request configuration. Shell commands start in `workspace_dir`, and `HOME` -is forced to `home_dir`. On Local, sibling materialized Homes may exist in the -same shellctl namespace, while path isolation restricts the active lease to its -own Home plus the shared Workspace. +is forced to `home_dir`. The standard temp variables `TMPDIR`, `TMP`, and `TEMP` +also point directly to `workspace_dir`, so the Workspace is both the command +`cwd` and temp space. On Local, sibling materialized Homes may exist in the same +shellctl namespace, while path isolation restricts the active lease to its own +Home plus the shared Workspace. ## Backend support diff --git a/dify-agent/docs/dify-agent/guide/index.md b/dify-agent/docs/dify-agent/guide/index.md index 3090dd6e0cb..99e2cc8e2b1 100644 --- a/dify-agent/docs/dify-agent/guide/index.md +++ b/dify-agent/docs/dify-agent/guide/index.md @@ -36,6 +36,7 @@ also reads `.env` and `dify-agent/.env` when present. | `DIFY_AGENT_SHUTDOWN_GRACE_SECONDS` | `30` | Seconds to wait for active local runs during graceful shutdown before cancellation. | | `DIFY_AGENT_RUN_RETENTION_SECONDS` | `259200` | Seconds to retain Redis run records and per-run event streams; defaults to 3 days. | | `DIFY_AGENT_RUN_TIMEOUT_SECONDS` | `3600` | Wall-clock deadline in seconds for the Pydantic AI `agent.run(...)` model/tool loop. Deadline failures use `agent_run_limit_exceeded`. Its default intentionally matches `DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS`, but the settings are independently configurable. | +| `DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS` | `210` | Shell command deadline for running the sandbox `dify-agent file upload --no-download-link` conversion. Keep it above the CLI's 180-second upload deadline. | | `DIFY_AGENT_API_TOKEN` | empty | Optional Bearer token required by private run, Execution Binding, Home Snapshot, and Binding file control-plane routes. Must match Dify API `AGENT_BACKEND_API_TOKEN`. | | `DIFY_AGENT_PLUGIN_DAEMON_URL` | `http://localhost:5002` | Base URL for the Dify plugin daemon. | | `DIFY_AGENT_PLUGIN_DAEMON_API_KEY` | empty | API key sent to the Dify plugin daemon. | @@ -44,9 +45,9 @@ also reads `.env` and `dify-agent/.env` when present. | `DIFY_AGENT_RUNTIME_BACKEND` | `local` | Selects one coherent `local`, `enterprise`, or `e2b` Home Snapshot + Execution Binding backend profile. | | `DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT` | empty | Local shellctl data-plane URL. With the default Local selection, leaving it empty disables `dify.runtime` and resource endpoints. | | `DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN` | empty | Optional bearer token sent to Local shellctl. | -| `DIFY_AGENT_LOCAL_SANDBOX_MATERIALIZED_HOME_ROOT` | `/home/dify/.dify-agent-materialized-homes` | Root directory, on the Local shellctl filesystem, for per-Binding materialized Homes. | -| `DIFY_AGENT_LOCAL_SANDBOX_WORKSPACE_ROOT` | `/home/dify/.dify-agent-workspaces` | Root directory, on the Local shellctl filesystem, for mutable Workspaces. | -| `DIFY_AGENT_LOCAL_SANDBOX_HOME_SNAPSHOT_ROOT` | `/home/dify/.dify-agent-home-snapshots` | Root directory, on the Local shellctl filesystem, for immutable Home Snapshots. | +| `DIFY_AGENT_LOCAL_SANDBOX_MATERIALIZED_HOME_ROOT` | `/home/dify` | Root directory, on the Local shellctl filesystem, for per-Binding materialized Homes. | +| `DIFY_AGENT_LOCAL_SANDBOX_WORKSPACE_ROOT` | `/workspace` | Root directory, on the Local shellctl filesystem, for mutable Workspaces. | +| `DIFY_AGENT_LOCAL_SANDBOX_HOME_SNAPSHOT_ROOT` | `/home/dify/.snapshots` | Root directory, on the Local shellctl filesystem, for immutable Home Snapshots. | | `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_ENDPOINT` | empty | Enterprise Gateway endpoint required by configuration. Default-Home Bindings are supported; immutable Home Snapshot operations remain unsupported. | | `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_AUTH_TOKEN` | empty | Optional `X-Inner-Api-Key` sent to the Enterprise Gateway. | | `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_TIMEOUT` | `30` | Enterprise control-plane timeout in seconds. | diff --git a/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md b/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md index 63f76ecce1d..b04d2f7360c 100644 --- a/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md @@ -257,7 +257,9 @@ The resource part serializes as: backend execution namespace. They are not host filesystem paths and are not sent in the run request. Shell commands start in `workspace_dir`, while `HOME` is forced to `home_dir`; `~` therefore resolves to the current Binding's -materialized Home. +materialized Home. The runner also sets `TMPDIR`, `TMP`, and `TEMP` directly to +`workspace_dir`, making the active Workspace both the default `cwd` and the +temporary working space. Workspace content persists with the Workspace until Dify API retires and collects it. Releasing a RuntimeLease ends only the current operation. Dify API can later diff --git a/dify-agent/src/dify_agent/adapters/shell/__init__.py b/dify-agent/src/dify_agent/adapters/shell/__init__.py index c09839b8d7b..6c5eb9b8399 100644 --- a/dify-agent/src/dify_agent/adapters/shell/__init__.py +++ b/dify-agent/src/dify_agent/adapters/shell/__init__.py @@ -9,6 +9,7 @@ from dify_agent.adapters.shell.protocols import ( ShellCommandProtocol, ShellCommandResult, ShellCommandStatus, + ShellExecutionMode, ShellPromptObservation, ShellProviderError, ) @@ -27,6 +28,7 @@ __all__ = [ "ShellCommandProtocol", "ShellCommandResult", "ShellCommandStatus", + "ShellExecutionMode", "ShellPromptObservation", "ShellProviderError", ] diff --git a/dify-agent/src/dify_agent/adapters/shell/protocols.py b/dify-agent/src/dify_agent/adapters/shell/protocols.py index 52b5436292a..d25da6a3443 100644 --- a/dify-agent/src/dify_agent/adapters/shell/protocols.py +++ b/dify-agent/src/dify_agent/adapters/shell/protocols.py @@ -4,6 +4,9 @@ from dataclasses import dataclass from typing import Literal, Protocol +type ShellExecutionMode = Literal["pty", "stdio"] + + @dataclass(frozen=True, slots=True) class ShellCommandResult: job_id: str @@ -63,6 +66,7 @@ class ShellCommandProtocol(Protocol): cwd: str | None = None, env: dict[str, str] | None = None, timeout: float, + mode: ShellExecutionMode = "pty", ) -> ShellCommandResult: ... async def wait( diff --git a/dify-agent/src/dify_agent/adapters/shell/shellctl.py b/dify-agent/src/dify_agent/adapters/shell/shellctl.py index cea953a4282..5af4e20afea 100644 --- a/dify-agent/src/dify_agent/adapters/shell/shellctl.py +++ b/dify-agent/src/dify_agent/adapters/shell/shellctl.py @@ -14,12 +14,13 @@ from typing import Protocol, TypeVar, cast import httpx2 as httpx from shellctl.client import ShellctlClientError -from shellctl.shared import HealthResponse +from shellctl.shared import HealthResponse, JobMode from dify_agent.adapters.shell.protocols import ( ShellCommandProtocol, ShellCommandResult, ShellCommandStatus, + ShellExecutionMode, ShellProviderError, ) @@ -60,6 +61,7 @@ class ShellctlClientProtocol(Protocol): cwd: str | None = None, env: dict[str, str] | None = None, timeout: float = _DEFAULT_TIMEOUT_SECONDS, + mode: JobMode = JobMode.PTY, ) -> ShellctlJobResult: ... async def wait( @@ -114,6 +116,7 @@ class ShellctlCommands(ShellCommandProtocol): cwd: str | None = None, env: dict[str, str] | None = None, timeout: float, + mode: ShellExecutionMode = "pty", ) -> ShellCommandResult: resolved_cwd = _resolve_lease_cwd( cwd, @@ -122,7 +125,15 @@ class ShellctlCommands(ShellCommandProtocol): ) resolved_env = _lease_env(env, home_dir=self.home_dir) return _from_job_result( - await _run_client_call(self.client.run(script, cwd=resolved_cwd, env=resolved_env, timeout=timeout)) + await _run_client_call( + self.client.run( + script, + cwd=resolved_cwd, + env=resolved_env, + timeout=timeout, + mode=JobMode(mode), + ) + ) ) async def wait( diff --git a/dify-agent/src/dify_agent/client/_client.py b/dify-agent/src/dify_agent/client/_client.py index 80af37d12fe..0594c0db497 100644 --- a/dify-agent/src/dify_agent/client/_client.py +++ b/dify-agent/src/dify_agent/client/_client.py @@ -52,7 +52,6 @@ from dify_agent.protocol import ( _ResponseModelT = TypeVar("_ResponseModelT", bound=BaseModel) _TERMINAL_EVENT_TYPES = {"run_succeeded", "run_failed", "run_cancelled"} _TERMINAL_RUN_STATUSES = {"succeeded", "failed", "cancelled"} -_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS = 90.0 _function_tool_result_payload_key_cache: str | None = None @@ -271,6 +270,7 @@ class Client: _base_url: str _timeout: float | httpx.Timeout _stream_timeout: float | httpx.Timeout | None + _binding_file_download_timeout: float | httpx.Timeout _headers: dict[str, str] _sync_http_client: httpx.Client | None _async_http_client: httpx.AsyncClient | None @@ -285,6 +285,7 @@ class Client: base_url: str, timeout: float | httpx.Timeout = 30.0, stream_timeout: float | httpx.Timeout | None = 30.0, + binding_file_download_timeout: float | httpx.Timeout = 240.0, headers: dict[str, str] | None = None, sync_http_client: httpx.Client | None = None, async_http_client: httpx.AsyncClient | None = None, @@ -292,6 +293,7 @@ class Client: self._base_url = base_url.rstrip("/") self._timeout = timeout self._stream_timeout = stream_timeout + self._binding_file_download_timeout = binding_file_download_timeout self._headers = dict(headers or {}) self._sync_http_client = sync_http_client self._async_http_client = async_http_client @@ -549,7 +551,7 @@ class Client: "download_binding_file", "/execution-bindings/files/download", request, - timeout=_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS, + timeout=self._binding_file_download_timeout, ) return _parse_model_response(response, BindingFileDownloadResponse) @@ -558,7 +560,7 @@ class Client: "download_binding_file_sync", "/execution-bindings/files/download", request, - timeout=_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS, + timeout=self._binding_file_download_timeout, ) return _parse_model_response(response, BindingFileDownloadResponse) diff --git a/dify-agent/src/dify_agent/layers/shell/layer.py b/dify-agent/src/dify_agent/layers/shell/layer.py index be7c73bdd56..2e39185aa64 100644 --- a/dify-agent/src/dify_agent/layers/shell/layer.py +++ b/dify-agent/src/dify_agent/layers/shell/layer.py @@ -108,9 +108,10 @@ Installed CLI: Filesystem spaces: -- `$HOME` is the system space. -- The current working directory (`cwd`) is the temporary working space. Relative paths resolve from here. -- Store temporary files under `/.tmp` (normally `./.tmp`). Do not use `/tmp`. +- `$HOME` is the system space for reusable tools and state. +- The current working directory (`cwd`) is the active Workspace and temporary working space. +- Relative paths and the standard temp environment variables (`TMPDIR`, `TMP`, and `TEMP`) resolve directly to `cwd`. +- Do not use `/tmp`. shell_run script rules: @@ -460,6 +461,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC ), timeout=timeout, max_output_bytes=max_output_bytes, + mode="stdio", ) async def run_remote_script( @@ -488,6 +490,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC env=self._build_shell_command_env(include_agent_stub_env=False), timeout=DEFAULT_TIMEOUT_SECONDS, max_output_bytes=_REMOTE_COMPLETE_OUTPUT_MAX_BYTES, + mode="stdio", ) def _require_resource(self) -> RuntimeLease: diff --git a/dify-agent/src/dify_agent/runtime/command_runner.py b/dify-agent/src/dify_agent/runtime/command_runner.py index 1ceec51a0cf..69af1f6e6b2 100644 --- a/dify-agent/src/dify_agent/runtime/command_runner.py +++ b/dify-agent/src/dify_agent/runtime/command_runner.py @@ -10,6 +10,7 @@ from dify_agent.adapters.shell.protocols import ( CompleteShellCommandResult, ShellCommandProtocol, ShellCommandResult, + ShellExecutionMode, ) from dify_agent.layers.shell.output_text import utf8_prefix @@ -26,6 +27,7 @@ async def execute_complete_with_commands( env: dict[str, str] | None, timeout: float, max_output_bytes: int, + mode: ShellExecutionMode, ) -> CompleteShellCommandResult: """Run a command to completion with bounded output and deterministic cleanup.""" @@ -36,7 +38,13 @@ async def execute_complete_with_commands( captured_bytes = 0 incomplete_reason: Literal["output_limit", "timeout"] | None = None try: - result = await commands.run(script, cwd=cwd, env=env, timeout=_remaining_time(deadline)) + result = await commands.run( + script, + cwd=cwd, + env=env, + timeout=_remaining_time(deadline), + mode=mode, + ) job_id = result.job_id while True: remaining_bytes = max(max_output_bytes - captured_bytes, 0) diff --git a/dify-agent/src/dify_agent/runtime_backend/e2b.py b/dify-agent/src/dify_agent/runtime_backend/e2b.py index 3140391eb13..4dad6ff13bc 100644 --- a/dify-agent/src/dify_agent/runtime_backend/e2b.py +++ b/dify-agent/src/dify_agent/runtime_backend/e2b.py @@ -49,11 +49,17 @@ class _E2BControlPlaneNotFoundError(RuntimeError): """Typed boundary error for SDK resources that no longer exist.""" +class _E2BFileEntry(Protocol): + path: str + + class _E2BFileSystem(Protocol): async def make_dir(self, path: str) -> bool: ... async def exists(self, path: str) -> bool: ... + async def list(self, path: str) -> list[_E2BFileEntry]: ... + async def remove(self, path: str) -> None: ... @@ -212,7 +218,7 @@ class E2BExecutionBindingBackend: active_timeout_seconds: int shellctl_port: int = 5004 layout: RuntimeLayout = field( - default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace") + default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/workspace") ) async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: @@ -233,9 +239,9 @@ class E2BExecutionBindingBackend: }, on_timeout="pause", ) - if await sandbox.files.exists(self.layout.workspace_dir): - await sandbox.files.remove(self.layout.workspace_dir) _ = await sandbox.files.make_dir(self.layout.workspace_dir) + for entry in await sandbox.files.list(self.layout.workspace_dir): + await sandbox.files.remove(entry.path) sandbox_id = sandbox.sandbox_id _ = await sandbox.pause(keep_memory=True) return ExecutionBindingAllocation(binding_ref=sandbox_id, workspace_ref=sandbox_id) diff --git a/dify-agent/src/dify_agent/runtime_backend/enterprise.py b/dify-agent/src/dify_agent/runtime_backend/enterprise.py index f86867d778e..c3675c3f6d4 100644 --- a/dify-agent/src/dify_agent/runtime_backend/enterprise.py +++ b/dify-agent/src/dify_agent/runtime_backend/enterprise.py @@ -70,7 +70,7 @@ class EnterpriseExecutionBindingBackend: gateway_timeout: float = 30.0 proxy_timeout: float = 60.0 layout: RuntimeLayout = field( - default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace") + default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/workspace") ) async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: @@ -104,8 +104,8 @@ class EnterpriseExecutionBindingBackend: [ "set -eu", f"mkdir -p {shlex.quote(self.layout.home_dir)}", - f"rm -rf -- {shlex.quote(self.layout.workspace_dir)}", f"mkdir -p {shlex.quote(self.layout.workspace_dir)}", + f"find {shlex.quote(self.layout.workspace_dir)} -mindepth 1 -maxdepth 1 -exec rm -rf -- {{}} +", f"chmod 700 {shlex.quote(self.layout.home_dir)} {shlex.quote(self.layout.workspace_dir)}", ] ), diff --git a/dify-agent/src/dify_agent/runtime_backend/local.py b/dify-agent/src/dify_agent/runtime_backend/local.py index f0a2e5c2e66..7ea883e2ec6 100644 --- a/dify-agent/src/dify_agent/runtime_backend/local.py +++ b/dify-agent/src/dify_agent/runtime_backend/local.py @@ -45,7 +45,7 @@ logger = logging.getLogger(__name__) class LocalHomeSnapshotBackend: endpoint: str auth_token: str - snapshot_root: str = "/home/dify/.dify-agent-home-snapshots" + snapshot_root: str = "/home/dify/.snapshots" client_factory: ShellctlClientFactory | None = None async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: @@ -111,9 +111,9 @@ class LocalHomeSnapshotBackend: class LocalExecutionBindingBackend: endpoint: str auth_token: str - materialized_home_root: str = "/home/dify/.dify-agent-materialized-homes" - workspace_root: str = "/home/dify/.dify-agent-workspaces" - snapshot_root: str = "/home/dify/.dify-agent-home-snapshots" + materialized_home_root: str = "/home/dify" + workspace_root: str = "/workspace" + snapshot_root: str = "/home/dify/.snapshots" client_factory: ShellctlClientFactory | None = None async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: @@ -238,9 +238,17 @@ class LocalExecutionBindingBackend: def _control_lease(self, handle: str) -> ShellctlRuntimeLease: control_root = _control_root((self.materialized_home_root, self.workspace_root, self.snapshot_root)) + layout = RuntimeLayout(home_dir=control_root, workspace_dir=control_root) + if control_root == "/": + # Keep the canonical root-level Workspace separate instead of + # broadening the control job's writable layout to the whole filesystem. + layout = RuntimeLayout( + home_dir=posixpath.commonpath((self.materialized_home_root, self.snapshot_root)), + workspace_dir=self.workspace_root, + ) return create_shellctl_lease( handle=handle, - layout=RuntimeLayout(home_dir=control_root, workspace_dir=control_root), + layout=layout, entrypoint=self.endpoint, token=self.auth_token, client_factory=self.client_factory, diff --git a/dify-agent/src/dify_agent/runtime_backend/profile.py b/dify-agent/src/dify_agent/runtime_backend/profile.py index 9bbbf3ee3f7..583afe54d4e 100644 --- a/dify-agent/src/dify_agent/runtime_backend/profile.py +++ b/dify-agent/src/dify_agent/runtime_backend/profile.py @@ -20,9 +20,9 @@ from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, Local from dify_agent.runtime_backend.protocols import RuntimeBackendProfile DEFAULT_E2B_TEMPLATE = "difys-default-team/dify-agent-local-sandbox" -DEFAULT_LOCAL_MATERIALIZED_HOME_ROOT = "/home/dify/.dify-agent-materialized-homes" -DEFAULT_LOCAL_WORKSPACE_ROOT = "/home/dify/.dify-agent-workspaces" -DEFAULT_LOCAL_HOME_SNAPSHOT_ROOT = "/home/dify/.dify-agent-home-snapshots" +DEFAULT_LOCAL_MATERIALIZED_HOME_ROOT = "/home/dify" +DEFAULT_LOCAL_WORKSPACE_ROOT = "/workspace" +DEFAULT_LOCAL_HOME_SNAPSHOT_ROOT = "/home/dify/.snapshots" class RuntimeBackendSettings(BaseSettings): diff --git a/dify-agent/src/dify_agent/runtime_backend/shellctl.py b/dify-agent/src/dify_agent/runtime_backend/shellctl.py index f947ea0d9ba..cd02ff1457c 100644 --- a/dify-agent/src/dify_agent/runtime_backend/shellctl.py +++ b/dify-agent/src/dify_agent/runtime_backend/shellctl.py @@ -112,8 +112,8 @@ async def run_shellctl_control_command( *, timeout: float = 30.0, ) -> CompleteShellCommandResult: - """Run one bounded driver control command and always delete its transient job.""" - result = await commands.run(script, cwd=None, env=None, timeout=timeout) + """Run one bounded control command through stdout-only stdio and delete its transient job.""" + result = await commands.run(script, cwd=None, env=None, timeout=timeout, mode="stdio") job_id = result.job_id output_parts = [result.output] try: diff --git a/dify-agent/src/dify_agent/server/app.py b/dify-agent/src/dify_agent/server/app.py index 0378e04baa5..8144c66ae20 100644 --- a/dify-agent/src/dify_agent/server/app.py +++ b/dify-agent/src/dify_agent/server/app.py @@ -76,6 +76,7 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI: execution_bindings=runtime_backend_profile.execution_bindings, agent_stub_api_base_url=resolved_settings.agent_stub_api_base_url, agent_stub_token_factory=agent_stub_token_factory, + download_command_timeout_seconds=resolved_settings.binding_file_download_command_timeout_seconds, ) if runtime_backend_profile is not None else None diff --git a/dify-agent/src/dify_agent/server/binding_files.py b/dify-agent/src/dify_agent/server/binding_files.py index acb24094eaf..a8d394bba3b 100644 --- a/dify-agent/src/dify_agent/server/binding_files.py +++ b/dify-agent/src/dify_agent/server/binding_files.py @@ -39,7 +39,6 @@ logger = logging.getLogger(__name__) _LIST_MAX_ENTRIES = 1000 _BROWSE_TIMEOUT_SECONDS = 60.0 _BROWSE_OUTPUT_MAX_BYTES = 1024 * 1024 -_DOWNLOAD_TIMEOUT_SECONDS = 60.0 _DOWNLOAD_OUTPUT_MAX_BYTES = 32 * 1024 _PAYLOAD_BEGIN = "<<>>" _PAYLOAD_END = "<<>>" @@ -149,6 +148,7 @@ class BindingFileService: execution_bindings: ExecutionBindingBackend agent_stub_api_base_url: str | None agent_stub_token_factory: ShellAgentStubTokenFactory | None + download_command_timeout_seconds: float async def list_files(self, request: BindingFileListRequest) -> BindingFileListResponse: try: @@ -166,6 +166,7 @@ class BindingFileService: env={"HOME": lease.layout.home_dir}, timeout=_BROWSE_TIMEOUT_SECONDS, max_output_bytes=_BROWSE_OUTPUT_MAX_BYTES, + mode="stdio", ) payload = _require_browse_payload(result, operation="list") try: @@ -193,6 +194,7 @@ class BindingFileService: env={"HOME": lease.layout.home_dir}, timeout=_BROWSE_TIMEOUT_SECONDS, max_output_bytes=_BROWSE_OUTPUT_MAX_BYTES, + mode="stdio", ) payload = _require_browse_payload(result, operation="read") try: @@ -241,8 +243,9 @@ class BindingFileService: f"dify-agent file upload --no-download-link {shlex.quote(resolved_path)}", cwd=lease.layout.workspace_dir, env=env, - timeout=_DOWNLOAD_TIMEOUT_SECONDS, + timeout=self.download_command_timeout_seconds, max_output_bytes=_DOWNLOAD_OUTPUT_MAX_BYTES, + mode="stdio", ) except ShellProviderError as exc: if exc.code == "timeout": diff --git a/dify-agent/src/dify_agent/server/settings.py b/dify-agent/src/dify_agent/server/settings.py index 866ce832d93..1601d3d019e 100644 --- a/dify-agent/src/dify_agent/server/settings.py +++ b/dify-agent/src/dify_agent/server/settings.py @@ -82,6 +82,7 @@ class ServerSettings(BaseSettings): description="Maximum Agent Stub upload size in MiB", validation_alias="DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT", ) + binding_file_download_command_timeout_seconds: float = Field(default=210.0, gt=0) server_secret_key: str | None = None api_token: str | None = None shell_redact_patterns: str = "" diff --git a/dify-agent/src/shellctl/__init__.py b/dify-agent/src/shellctl/__init__.py index cb294b365cf..4f7b153c5c8 100644 --- a/dify-agent/src/shellctl/__init__.py +++ b/dify-agent/src/shellctl/__init__.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: HealthResponse, InputJobRequest, JobInfo, + JobMode, JobResult, JobStatusName, JobStatusView, @@ -63,6 +64,7 @@ __all__ = [ "HealthResponse", "InputJobRequest", "JobInfo", + "JobMode", "JobResult", "JobStatusName", "JobStatusView", @@ -97,6 +99,7 @@ _EXPORTS = { "HealthResponse": "shellctl.shared", "InputJobRequest": "shellctl.shared", "JobInfo": "shellctl.shared", + "JobMode": "shellctl.shared", "JobResult": "shellctl.shared", "JobStatusName": "shellctl.shared", "JobStatusView": "shellctl.shared", diff --git a/dify-agent/src/shellctl/client/sdk.py b/dify-agent/src/shellctl/client/sdk.py index cc8ee05ca9f..ebac8135d21 100644 --- a/dify-agent/src/shellctl/client/sdk.py +++ b/dify-agent/src/shellctl/client/sdk.py @@ -29,6 +29,7 @@ from shellctl.shared.schemas import ( DeleteJobResponse, HealthResponse, JobInfo, + JobMode, JobResult, JobStatusView, ListJobsResponse, @@ -144,11 +145,14 @@ class ShellctlClient: env: dict[str, str] | None = None, timeout: float = DEFAULT_TIMEOUT_SECONDS, terminal: TerminalSize | None = None, + mode: JobMode = JobMode.PTY, ) -> JobResult: """Create a new job and wait for initial output or completion. `cwd` and `env` preset the script's working directory and environment - overlay on the server side. + overlay on the server side. `mode="stdio"` provides stdout-only public + output for non-interactive commands; the default PTY mode remains + interactive and merges stdout with stderr. """ payload = RunJobRequest( @@ -156,6 +160,7 @@ class ShellctlClient: cwd=cwd, env=env, terminal=terminal, + mode=mode, timeout=timeout, output_limit=self.output_limit, idle_flush_seconds=self.idle_flush_seconds, diff --git a/dify-agent/src/shellctl/shared/__init__.py b/dify-agent/src/shellctl/shared/__init__.py index e53a46d5e27..bb79cb98742 100644 --- a/dify-agent/src/shellctl/shared/__init__.py +++ b/dify-agent/src/shellctl/shared/__init__.py @@ -60,6 +60,7 @@ if TYPE_CHECKING: HealthResponse, InputJobRequest, JobInfo, + JobMode, JobResult, JobStatusName, JobStatusView, @@ -101,6 +102,7 @@ __all__ = [ "HealthResponse", "InputJobRequest", "JobInfo", + "JobMode", "JobResult", "JobStatusName", "JobStatusView", @@ -166,6 +168,7 @@ _EXPORTS = { "HealthResponse": "shellctl.shared.schemas", "InputJobRequest": "shellctl.shared.schemas", "JobInfo": "shellctl.shared.schemas", + "JobMode": "shellctl.shared.schemas", "JobResult": "shellctl.shared.schemas", "JobStatusName": "shellctl.shared.schemas", "JobStatusView": "shellctl.shared.schemas", diff --git a/dify-agent/src/shellctl/shared/schemas.py b/dify-agent/src/shellctl/shared/schemas.py index 0bbaf323a18..7c3e5df7a0c 100644 --- a/dify-agent/src/shellctl/shared/schemas.py +++ b/dify-agent/src/shellctl/shared/schemas.py @@ -42,6 +42,13 @@ class JobStatusName(StrEnum): LOST = "lost" +class JobMode(StrEnum): + """Standard-stream wiring used to execute a shellctl job.""" + + PTY = "pty" + STDIO = "stdio" + + TERMINAL_JOB_STATUSES = frozenset( { JobStatusName.EXITED, @@ -139,6 +146,7 @@ class RunJobRequest(ShellctlModel): cwd: str | None = None env: dict[str, str] | None = None terminal: TerminalSize | None = None + mode: JobMode = JobMode.PTY timeout: float = Field(default=DEFAULT_TIMEOUT_SECONDS, gt=0, le=SHELL_TOOL_HARD_TIMEOUT_SECONDS) output_limit: int = Field(default=DEFAULT_OUTPUT_LIMIT_BYTES, ge=1, le=MAX_OUTPUT_LIMIT_BYTES) idle_flush_seconds: float = Field(default=DEFAULT_IDLE_FLUSH_SECONDS, ge=0, le=30) @@ -201,6 +209,7 @@ __all__ = [ "HealthResponse", "InputJobRequest", "JobInfo", + "JobMode", "JobResult", "JobStatusName", "JobStatusView", diff --git a/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py b/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py index 6984d3946ad..d2d8b308942 100644 --- a/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py +++ b/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py @@ -9,6 +9,7 @@ from typing import cast import httpx2 as httpx import pytest from shellctl.client import ShellctlClientError +from shellctl.shared import JobMode from dify_agent.adapters.shell.protocols import ShellProviderError from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol, ShellctlCommands @@ -39,12 +40,20 @@ class _Status: class _Client: run_result: object = field(default_factory=_Job) delete_error: Exception | None = None - run_calls: list[tuple[str, str | None, dict[str, str] | None, float]] = field(default_factory=list) + run_calls: list[tuple[str, str | None, dict[str, str] | None, float, JobMode]] = field(default_factory=list) wait_calls: list[tuple[str, int, float]] = field(default_factory=list) delete_calls: list[tuple[str, bool, float | None]] = field(default_factory=list) - async def run(self, script: str, *, cwd=None, env=None, timeout=30.0): - self.run_calls.append((script, cwd, env, timeout)) + async def run( + self, + script: str, + *, + cwd=None, + env=None, + timeout=30.0, + mode: JobMode = JobMode.PTY, + ): + self.run_calls.append((script, cwd, env, timeout, mode)) if isinstance(self.run_result, Exception): raise self.run_result return self.run_result @@ -85,7 +94,20 @@ def test_commands_apply_runtime_layout_and_home_environment() -> None: assert result.output == "ok" asyncio.run(scenario()) - assert client.run_calls == [("pwd", "/workspace/reports", {"TOKEN": "value", "HOME": "/home/binding"}, 2.5)] + assert client.run_calls == [ + ("pwd", "/workspace/reports", {"TOKEN": "value", "HOME": "/home/binding"}, 2.5, JobMode.PTY) + ] + + +def test_commands_forward_stdio_mode() -> None: + client = _Client() + + async def scenario() -> None: + commands = ShellctlCommands(_client(client)) + await commands.run("printf result", timeout=2.5, mode="stdio") + + asyncio.run(scenario()) + assert client.run_calls == [("printf result", None, None, 2.5, JobMode.STDIO)] def test_commands_reject_cwd_outside_runtime_layout() -> None: diff --git a/dify-agent/tests/local/dify_agent/client/test_client.py b/dify-agent/tests/local/dify_agent/client/test_client.py index 54c905625ac..3f2a4001786 100644 --- a/dify-agent/tests/local/dify_agent/client/test_client.py +++ b/dify-agent/tests/local/dify_agent/client/test_client.py @@ -98,9 +98,9 @@ def _binding_file_download_request(path: str = "report.txt") -> BindingFileDownl ) -def _assert_binding_download_timeout(request: httpx.Request) -> None: +def _assert_binding_download_timeout(request: httpx.Request, expected: float = 240.0) -> None: timeout = cast(dict[str, float], request.extensions["timeout"]) - assert timeout == {"connect": 90.0, "read": 90.0, "write": 90.0, "pool": 90.0} + assert timeout == {"connect": expected, "read": expected, "write": expected, "pool": expected} def _function_tool_result_payload(key: str) -> dict[str, object]: @@ -398,13 +398,17 @@ def test_async_binding_file_methods_post_dtos_and_parse_responses() -> None: 200, json={"path": "note.txt", "size": 5, "truncated": False, "binary": False, "text": "hello"} ) if request.url.path == "/execution-bindings/files/download": - _assert_binding_download_timeout(request) + _assert_binding_download_timeout(request, expected=123.5) return httpx.Response(200, json={"reference": "dify-file-ref:file-1"}) raise AssertionError(f"unexpected request: {request.method} {request.url}") async def scenario() -> None: http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) - client = Client(base_url="http://testserver", async_http_client=http_client) + client = Client( + base_url="http://testserver", + binding_file_download_timeout=123.5, + async_http_client=http_client, + ) listing = await client.list_binding_files("binding-ref", ".") preview = await client.read_binding_file("binding-ref", "note.txt") diff --git a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py index b3362d2c033..0ffc47c9bd8 100644 --- a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py @@ -26,6 +26,7 @@ from dify_agent.layers.shell.layer import ( from dify_agent.adapters.shell.protocols import ( ShellCommandResult, ShellCommandStatus, + ShellExecutionMode, ShellProviderError, ) from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig @@ -113,6 +114,7 @@ class RunCall: cwd: str | None env: Mapping[str, str] | None timeout: float + mode: ShellExecutionMode = "pty" @dataclass(slots=True) @@ -167,8 +169,16 @@ class FakeCommands: interrupt_calls: list[InterruptCall] = field(default_factory=list) delete_calls: list[DeleteCall] = field(default_factory=list) - async def run(self, script: str, *, cwd: str | None = None, env: dict[str, str] | None = None, timeout: float): - self.run_calls.append(RunCall(script=script, cwd=cwd, env=env, timeout=timeout)) + async def run( + self, + script: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout: float, + mode: ShellExecutionMode = "pty", + ): + self.run_calls.append(RunCall(script=script, cwd=cwd, env=env, timeout=timeout, mode=mode)) if self.run_handler is None: raise AssertionError("Unexpected run() call") return self.run_handler(script, cwd, env, timeout) @@ -293,6 +303,15 @@ def test_shell_type_id_constant_matches_implementation_class() -> None: assert DIFY_SHELL_LAYER_TYPE_ID == DifyShellLayer.type_id +def test_shell_prefix_prompt_describes_workspace_as_temp_space() -> None: + prompt = shell_layer_module._SHELL_LAYER_PREFIX_PROMPT + + assert "`cwd`) is the active Workspace and temporary working space" in prompt + assert "`TMPDIR`, `TMP`, and `TEMP`) resolve directly to `cwd`" in prompt + assert "`$HOME` is the system space for reusable tools and state" in prompt + assert "/.tmp" not in prompt + + def test_shell_layer_create_bootstraps_inside_sandbox_workspace() -> None: expected_home = "/home/agent-1" expected_workspace_cwd = "/home/agent-1/workspace/abc12ff" @@ -499,6 +518,7 @@ def test_shell_layer_tools_map_inputs_and_maintain_offsets_with_tail_end() -> No asyncio.run(scenario()) assert layer.runtime_state.job_offsets == {"user-job": 34} + assert commands.run_calls[0].mode == "pty" assert commands.tail_calls == [TailCall(job_id="user-job"), TailCall(job_id="user-job")] @@ -936,6 +956,7 @@ def test_run_remote_script_complete_uses_read_output_before_wait_and_deletes_job asyncio.run(scenario()) assert events == ["run", "read_output", "wait"] + assert commands.run_calls[0].mode == "stdio" assert [call.job_id for call in commands.delete_calls] == ["remote-job"] diff --git a/dify-agent/tests/local/dify_agent/runtime/test_command_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_command_runner.py index 2a8e9f6b658..ebe7b5ecea8 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_command_runner.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_command_runner.py @@ -6,7 +6,7 @@ from dataclasses import dataclass, field import pytest -from dify_agent.adapters.shell.protocols import ShellCommandResult +from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellExecutionMode from dify_agent.runtime.command_runner import execute_complete_with_commands @@ -16,11 +16,20 @@ class _BlockingCommands: wait_forever: asyncio.Event = field(default_factory=asyncio.Event) deletes: list[tuple[str, bool]] = field(default_factory=list) - async def run(self, script: str, *, cwd: str | None, env: dict[str, str] | None, timeout: float): + async def run( + self, + script: str, + *, + cwd: str | None, + env: dict[str, str] | None, + timeout: float, + mode: ShellExecutionMode = "pty", + ): assert script == "long-running" assert cwd == "/workspace" assert env == {"HOME": "/home/agent"} assert timeout > 0 + assert mode == "stdio" return ShellCommandResult( job_id="job-1", status="running", @@ -66,6 +75,7 @@ async def test_cancellation_deletes_job_returned_before_blocking_wait() -> None: env={"HOME": "/home/agent"}, timeout=60.0, max_output_bytes=4096, + mode="stdio", ) ) try: diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py index 9dcddff2b8c..62963e7066a 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py @@ -1,5 +1,6 @@ from __future__ import annotations +import posixpath from collections.abc import Callable from dataclasses import dataclass, field from typing import cast @@ -28,9 +29,15 @@ from dify_agent.runtime_backend.e2b import ( from dify_agent.runtime_backend.shellctl import ShellctlRuntimeLease +@dataclass(frozen=True, slots=True) +class _FileEntry: + path: str + + @dataclass(slots=True) class _Files: paths: set[str] = field(default_factory=set) + removed: list[str] = field(default_factory=list) async def make_dir(self, path: str) -> bool: self.paths.add(path) @@ -39,8 +46,17 @@ class _Files: async def exists(self, path: str) -> bool: return path in self.paths + async def list(self, path: str) -> list[_FileEntry]: + prefix = f"{path.rstrip('/')}/" + return [ + _FileEntry(path=entry) + for entry in sorted(self.paths) + if entry.startswith(prefix) and "/" not in entry.removeprefix(prefix) + ] + async def remove(self, path: str) -> None: - self.paths.discard(path) + self.removed.append(path) + self.paths = {entry for entry in self.paths if entry != path and posixpath.commonpath((entry, path)) != path} @dataclass(slots=True) @@ -90,6 +106,14 @@ class _ControlPlane: del timeout sandbox_id = f"sandbox-{len(self.sandboxes) + 1}" sandbox = _Sandbox(sandbox_id=sandbox_id, pause_error=self.pause_error) + sandbox.files.paths.update( + { + "/workspace", + "/workspace/stale-dir", + "/workspace/stale-dir/nested.txt", + "/workspace/stale.txt", + } + ) self.sandboxes[sandbox_id] = sandbox self.created.append((template, on_timeout)) assert metadata["dify.resource"] == "runtime-sandbox" @@ -129,7 +153,7 @@ def _mock_http( def _connected_backend(*, pause_error: Exception | None = None) -> tuple[E2BExecutionBindingBackend, _Sandbox]: control = _ControlPlane() sandbox = _Sandbox(sandbox_id="sandbox-1", pause_error=pause_error) - sandbox.files.paths.add("/home/dify/workspace") + sandbox.files.paths.add("/workspace") control.sandboxes[sandbox.sandbox_id] = sandbox return ( E2BExecutionBindingBackend( @@ -207,9 +231,12 @@ async def test_e2b_binding_uses_default_template_or_exact_snapshot_and_couples_r assert control.created == [("prepared-template", "pause"), ("snapshot-1", "pause")] assert default_allocation.binding_ref == default_allocation.workspace_ref assert snapshot_allocation.binding_ref == snapshot_allocation.workspace_ref - runtime = control.sandboxes[default_allocation.binding_ref] - assert runtime.files.paths == {"/home/dify/workspace"} - assert runtime.pauses == [True] + assert control.sandboxes[default_allocation.binding_ref].pauses == [True] + + for allocation in (default_allocation, snapshot_allocation): + runtime = control.sandboxes[allocation.binding_ref] + assert runtime.files.paths == {"/workspace"} + assert "/workspace" not in runtime.files.removed for allocation in (default_allocation, snapshot_allocation): await bindings.destroy_binding( @@ -358,6 +385,8 @@ async def test_e2b_acquire_retries_transient_shellctl_failures_until_ready( assert attempts == 3 assert sleeps == [0.5, 0.5] + assert lease.layout.home_dir == "/home/dify" + assert lease.layout.workspace_dir == "/workspace" assert not clients[0].is_closed await backend.release(lease) assert clients[0].is_closed diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py index 841e6d63a91..049d6fb65e6 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py @@ -80,7 +80,7 @@ async def test_enterprise_acquire_exposes_canonical_layout_through_gateway_proxy script = payload["script"] assert isinstance(script, str) assert "test -d /home/dify" in script - assert "test -d /home/dify/workspace" in script + assert "test -d /workspace" in script return _job_response() return httpx.Response(200, json={"job_id": "job-1"}) @@ -94,7 +94,7 @@ async def test_enterprise_acquire_exposes_canonical_layout_through_gateway_proxy lease = await backend.acquire("sandbox-1") assert lease.layout.home_dir == "/home/dify" - assert lease.layout.workspace_dir == "/home/dify/workspace" + assert lease.layout.workspace_dir == "/workspace" assert [request.url.path for request in requests] == [ "/proxy/v1/jobs/run", "/proxy/v1/jobs/job-1", @@ -299,7 +299,8 @@ async def test_enterprise_default_binding_creates_gateway_sandbox_and_layout( script = payload["script"] assert isinstance(script, str) assert "mkdir -p /home/dify" in script - assert "rm -rf -- /home/dify/workspace" in script + assert "mkdir -p /workspace" in script + assert "find /workspace -mindepth 1 -maxdepth 1 -exec rm -rf -- {} +" in script return _job_response() return httpx.Response(200, json={"job_id": "job-1"}) diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py index a2f8190e32f..5127ebd3c69 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py @@ -5,7 +5,7 @@ import shlex from typing import Mapping import pytest -from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView +from shellctl.shared import DeleteJobResponse, JobMode, JobResult, JobStatusName, JobStatusView from dify_agent.runtime_backend import ( BindingCreateError, @@ -22,6 +22,7 @@ class _RunCall: commands: tuple[tuple[str, ...], ...] cwd: str | None env: Mapping[str, str] | None + mode: JobMode @dataclass(slots=True) @@ -40,12 +41,13 @@ class _Client: cwd: str | None = None, env: Mapping[str, str] | None = None, timeout: float = 10.0, + mode: JobMode = JobMode.PTY, ) -> JobResult: del timeout commands = tuple( tuple(shlex.split(line)) for line in script.splitlines() if line.strip() and line.strip() != "set -eu" ) - self.runs.append(_RunCall(commands=commands, cwd=cwd, env=env)) + self.runs.append(_RunCall(commands=commands, cwd=cwd, env=env, mode=mode)) return JobResult( job_id=f"job-{len(self.runs)}", status=JobStatusName.EXITED, @@ -165,6 +167,7 @@ async def test_local_binding_create_materializes_home_and_new_workspace() -> Non assert ("mkdir", "-p", "/homes/binding-1") in factory.commands assert ("cp", "-a", "/snapshots/home-home-1/.", "/homes/binding-1/") in factory.commands assert ("chmod", "700", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands + assert all(run.mode is JobMode.STDIO for run in factory.runs) @pytest.mark.anyio @@ -195,6 +198,57 @@ async def test_local_binding_create_uses_empty_default_home_without_snapshot_acc assert all("/snapshots" not in part for command in factory.commands for part in command) +@pytest.mark.anyio +async def test_local_binding_bootstraps_custom_roots_from_common_root() -> None: + factory = _Factory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + materialized_home_root="/tmp/dify-agent/homes", + workspace_root="/tmp/dify-agent/workspaces", + snapshot_root="/tmp/dify-agent/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref=None, + ) + ) + + assert factory.runs[0].cwd == "/tmp/dify-agent" + assert factory.runs[0].env == {"HOME": "/tmp/dify-agent"} + + +@pytest.mark.anyio +async def test_local_binding_separates_root_workspace_from_home_control_scope() -> None: + factory = _Factory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref=None, + ) + ) + + assert factory.runs[0].cwd == "/workspace" + assert factory.runs[0].env == {"HOME": "/home/dify"} + + @pytest.mark.anyio async def test_local_binding_create_failure_removes_partial_home_and_workspace() -> None: factory = _FailThenSucceedFactory() @@ -245,6 +299,7 @@ async def test_local_binding_acquire_scopes_commands_to_materialized_home_and_wo pwd_run = next(run for run in factory.runs if run.commands == (("pwd",),)) assert pwd_run.cwd == "/workspaces/workspace-1" assert pwd_run.env == {"HOME": "/homes/binding-1"} + assert pwd_run.mode is JobMode.PTY with pytest.raises(ValueError, match="outside this RuntimeLease"): await lease.commands.run("cat secret", cwd="/homes/other", timeout=10.0) await backend.release(lease) @@ -320,87 +375,6 @@ async def test_local_snapshot_delete_removes_snapshot_directory() -> None: assert ("rm", "-rf", "--", "/snapshots/home-home-2") in factory.commands -@pytest.mark.anyio -async def test_local_backend_materializes_same_agent_twice_in_one_workspace() -> None: - factory = _Factory() - backend = LocalExecutionBindingBackend( - endpoint="http://shellctl", - auth_token="", - client_factory=factory, # pyright: ignore[reportArgumentType] - ) - - first = await backend.create_binding( - ExecutionBindingCreateSpec( - tenant_id="tenant-1", - agent_id="agent-1", - binding_id="binding-1", - workspace_id="workspace-1", - existing_workspace_ref=None, - home_snapshot_ref="home-home-1", - ) - ) - second = await backend.create_binding( - ExecutionBindingCreateSpec( - tenant_id="tenant-1", - agent_id="agent-1", - binding_id="binding-2", - workspace_id="workspace-1", - existing_workspace_ref=first.workspace_ref, - home_snapshot_ref="home-home-1", - ) - ) - - assert first.binding_ref == "binding-1:workspace-1" - assert second.binding_ref == "binding-2:workspace-1" - assert first.workspace_ref == second.workspace_ref == "workspace-1" - first_lease = await backend.acquire(first.binding_ref) - second_lease = await backend.acquire(second.binding_ref) - assert first_lease.layout.home_dir != second_lease.layout.home_dir - assert first_lease.layout.workspace_dir == second_lease.layout.workspace_dir - await backend.release(first_lease) - await backend.release(second_lease) - - await backend.destroy_binding( - ExecutionBindingDestroySpec( - binding_ref=first.binding_ref, - destroy_workspace=False, - ) - ) - surviving_lease = await backend.acquire(second.binding_ref) - assert surviving_lease.layout.home_dir == "/home/dify/.dify-agent-materialized-homes/binding-2" - assert surviving_lease.layout.workspace_dir == "/home/dify/.dify-agent-workspaces/workspace-1" - await backend.release(surviving_lease) - - workspace_dir = "/home/dify/.dify-agent-workspaces/workspace-1" - assert ("test", "-d", workspace_dir) in factory.commands - assert factory.commands.count(("mkdir", "-p", workspace_dir)) == 1 - assert ( - "rm", - "-rf", - "--", - "/home/dify/.dify-agent-materialized-homes/binding-1", - ) in factory.commands - assert ( - "rm", - "-rf", - "--", - "/home/dify/.dify-agent-materialized-homes/binding-1", - workspace_dir, - ) not in factory.commands - assert ( - "cp", - "-a", - "/home/dify/.dify-agent-home-snapshots/home-home-1/.", - "/home/dify/.dify-agent-materialized-homes/binding-1/", - ) in factory.commands - assert ( - "cp", - "-a", - "/home/dify/.dify-agent-home-snapshots/home-home-1/.", - "/home/dify/.dify-agent-materialized-homes/binding-2/", - ) in factory.commands - - @pytest.mark.anyio async def test_local_snapshot_delete_preserves_shellctl_error_when_close_fails() -> None: factory = _FailingFactory() diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py index fe6704ac969..1ae8b70005a 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py @@ -40,6 +40,19 @@ def test_local_backend_requires_shellctl_endpoint() -> None: _ = RuntimeBackendSettings(runtime_backend="local") +def test_local_backend_uses_root_workspace_directory_by_default() -> None: + settings = RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="http://shellctl.example", + ) + + profile = create_runtime_backend_profile(settings) + + assert settings.local_sandbox_workspace_root == "/workspace" + assert isinstance(profile.execution_bindings, LocalExecutionBindingBackend) + assert profile.execution_bindings.workspace_root == "/workspace" + + def test_local_backend_passes_configured_roots_to_drivers() -> None: settings = RuntimeBackendSettings( runtime_backend="local", diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py index 9d07ccfb027..ab94fa4ece1 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py @@ -5,7 +5,7 @@ from typing import cast import pytest -from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellCommandStatus +from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellCommandStatus, ShellExecutionMode from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol from dify_agent.runtime_backend.protocols import RuntimeLayout from dify_agent.runtime_backend.shellctl import ( @@ -43,6 +43,7 @@ class _FakeCommands: wait_error: Exception | None = None delete_error: Exception | None = None delete_calls: list[tuple[str, bool]] = field(default_factory=list) + run_modes: list[ShellExecutionMode] = field(default_factory=list) async def run( self, @@ -51,8 +52,10 @@ class _FakeCommands: cwd: str | None = None, env: dict[str, str] | None = None, timeout: float, + mode: ShellExecutionMode = "pty", ) -> ShellCommandResult: del script, cwd, env, timeout + self.run_modes.append(mode) return self.initial async def wait(self, job_id: str, *, offset: int, timeout: float) -> ShellCommandResult: @@ -169,6 +172,7 @@ async def test_control_command_success_is_preserved_when_delete_fails( result = await run_shellctl_control_command(commands, "true") assert result.output == "ok" + assert commands.run_modes == ["stdio"] assert commands.delete_calls == [("job-1", True)] assert "delete failed" in caplog.text @@ -188,4 +192,5 @@ async def test_control_command_error_is_preserved_when_delete_also_fails( _ = await run_shellctl_control_command(commands, "false") assert commands.delete_calls == [("job-1", True)] + assert commands.run_modes == ["stdio"] assert "delete failed" in caplog.text diff --git a/dify-agent/tests/local/dify_agent/server/test_binding_files.py b/dify-agent/tests/local/dify_agent/server/test_binding_files.py index 4975729a9dd..68c7f38a3e0 100644 --- a/dify-agent/tests/local/dify_agent/server/test_binding_files.py +++ b/dify-agent/tests/local/dify_agent/server/test_binding_files.py @@ -15,7 +15,7 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellProviderError +from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellExecutionMode, ShellProviderError from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.protocol import BindingFileDownloadRequest, BindingFileListRequest, BindingFileReadRequest from dify_agent.runtime_backend import BindingAcquireError, BindingLostError, RuntimeLayout, RuntimeLease @@ -38,7 +38,16 @@ class _Commands: calls: list[tuple[str, str | None, dict[str, str] | None, float]] = field(default_factory=list) deletes: list[str] = field(default_factory=list) - async def run(self, script: str, *, cwd: str | None = None, env=None, timeout: float) -> ShellCommandResult: + async def run( + self, + script: str, + *, + cwd: str | None = None, + env=None, + timeout: float, + mode: ShellExecutionMode = "pty", + ) -> ShellCommandResult: + assert mode == "stdio" self.calls.append((script, cwd, env, timeout)) output = self.outputs.pop(0) exit_code = self.exit_codes.pop(0) if self.exit_codes else 0 @@ -77,7 +86,16 @@ class _ProviderErrorCommands(_Commands): phase: Literal["run", "wait"] = "run" error_code: str = "timeout" - async def run(self, script: str, *, cwd: str | None = None, env=None, timeout: float) -> ShellCommandResult: + async def run( + self, + script: str, + *, + cwd: str | None = None, + env=None, + timeout: float, + mode: ShellExecutionMode = "pty", + ) -> ShellCommandResult: + assert mode == "stdio" self.calls.append((script, cwd, env, timeout)) if self.phase == "run": raise ShellProviderError("shell provider failed", code=self.error_code) @@ -99,7 +117,16 @@ class _ProviderErrorCommands(_Commands): @dataclass(slots=True) class _LocalCommands(_Commands): - async def run(self, script: str, *, cwd: str | None = None, env=None, timeout: float) -> ShellCommandResult: + async def run( + self, + script: str, + *, + cwd: str | None = None, + env=None, + timeout: float, + mode: ShellExecutionMode = "pty", + ) -> ShellCommandResult: + assert mode == "stdio" self.calls.append((script, cwd, env, timeout)) process = await asyncio.create_subprocess_shell( script, @@ -152,12 +179,18 @@ def _context() -> DifyExecutionContextLayerConfig: ) -def _service(commands: _Commands, *, configured: bool = True) -> tuple[BindingFileService, _Backend]: +def _service( + commands: _Commands, + *, + configured: bool = True, + download_command_timeout_seconds: float = 210.0, +) -> tuple[BindingFileService, _Backend]: backend = _Backend(lease=cast(RuntimeLease, _Lease(commands=commands))) service = BindingFileService( execution_bindings=backend, # pyright: ignore[reportArgumentType] agent_stub_api_base_url="http://stub/agent-stub" if configured else None, agent_stub_token_factory=(lambda execution_context, *, session_id: "secret-jwe") if configured else None, + download_command_timeout_seconds=download_command_timeout_seconds, ) return service, backend @@ -181,6 +214,7 @@ def _local_service(tmp_path: Path) -> tuple[BindingFileService, _Backend, _Local execution_bindings=backend, # pyright: ignore[reportArgumentType] agent_stub_api_base_url=None, agent_stub_token_factory=None, + download_command_timeout_seconds=210.0, ) return service, backend, commands, workspace, home @@ -388,7 +422,7 @@ async def test_download_shell_quotes_resolved_path_and_returns_only_reference_in commands = _Commands( outputs=[json.dumps({"transfer_method": "tool_file", "reference": _REFERENCE, "public_download_url": "bad"})] ) - service, backend = _service(commands) + service, backend = _service(commands, download_command_timeout_seconds=123.5) issued_tokens: list[tuple[DifyExecutionContextLayerConfig, str | None]] = [] def issue_token(execution_context: DifyExecutionContextLayerConfig, *, session_id: str | None) -> str: @@ -421,7 +455,7 @@ async def test_download_shell_quotes_resolved_path_and_returns_only_reference_in "DIFY_AGENT_STUB_API_BASE_URL": "http://stub/agent-stub", "DIFY_AGENT_STUB_AUTH_JWE": "secret-jwe", } - assert timeout == pytest.approx(60.0, rel=0, abs=0.01) + assert timeout == pytest.approx(123.5, rel=0, abs=0.01) assert issued_tokens == [(context, None)] assert issued_tokens[0][0].model_dump() == context.model_dump() assert backend.releases == 1 @@ -583,6 +617,7 @@ async def test_download_maps_binding_acquire_errors( execution_bindings=backend, # pyright: ignore[reportArgumentType] agent_stub_api_base_url="http://stub/agent-stub", agent_stub_token_factory=lambda execution_context, *, session_id: "secret-jwe", + download_command_timeout_seconds=210.0, ) with pytest.raises(BindingFileError) as exc_info: diff --git a/dify-agent/tests/local/dify_agent/server/test_settings.py b/dify-agent/tests/local/dify_agent/server/test_settings.py index c04d87fa764..76cdb54bec4 100644 --- a/dify-agent/tests/local/dify_agent/server/test_settings.py +++ b/dify-agent/tests/local/dify_agent/server/test_settings.py @@ -79,6 +79,31 @@ def test_server_settings_rejects_non_positive_run_timeout() -> None: _ = ServerSettings(run_timeout_seconds=0) +def test_server_settings_reads_binding_file_download_command_timeout_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS", "123.5") + + settings = ServerSettings() + + assert settings.binding_file_download_command_timeout_seconds == 123.5 + + +def test_server_settings_defaults_binding_file_download_command_timeout_to_210_seconds( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.delenv("DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS", raising=False) + monkeypatch.chdir(tmp_path) + + assert ServerSettings().binding_file_download_command_timeout_seconds == 210.0 + + +def test_server_settings_rejects_non_positive_binding_file_download_command_timeout() -> None: + with pytest.raises(ValidationError, match="greater than 0"): + _ = ServerSettings(binding_file_download_command_timeout_seconds=0) + + def test_server_settings_defaults_shellctl_auth_token_to_none( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/dify-agent/tests/local/shellctl/test_shellctl_client.py b/dify-agent/tests/local/shellctl/test_shellctl_client.py index fad5d147a27..8b09ed3a587 100644 --- a/dify-agent/tests/local/shellctl/test_shellctl_client.py +++ b/dify-agent/tests/local/shellctl/test_shellctl_client.py @@ -10,6 +10,7 @@ from shellctl.client import sdk as shellctl_sdk from shellctl.shared import ( DEFAULT_TERMINATE_GRACE_SECONDS, HealthResponse, + JobMode, JobStatusName, ) @@ -29,12 +30,14 @@ class ForcedDeleteKwargs(TypedDict, total=False): cwd="/tmp", env={"HELLO": "world"}, timeout=12, + mode=JobMode.STDIO, ), "/v1/jobs/run", { "script": "printf ready\\n", "cwd": "/tmp", "env": {"HELLO": "world"}, + "mode": "stdio", "timeout": 12.0, "output_limit": 4096, "idle_flush_seconds": 0.25, diff --git a/dify-agent/tests/local/shellctl/test_shellctl_shared.py b/dify-agent/tests/local/shellctl/test_shellctl_shared.py index 38fd14987c4..06fdd607d36 100644 --- a/dify-agent/tests/local/shellctl/test_shellctl_shared.py +++ b/dify-agent/tests/local/shellctl/test_shellctl_shared.py @@ -7,6 +7,7 @@ from pydantic import ValidationError from shellctl.shared import ( JOB_ID_ALPHABET, + JobMode, MAX_WAIT_TIMEOUT_SECONDS, RunJobRequest, SHELL_TOOL_HARD_TIMEOUT_SECONDS, @@ -18,6 +19,13 @@ from shellctl.shared import ( ) +def test_run_job_request_defaults_to_pty_and_rejects_unknown_mode() -> None: + assert RunJobRequest(script="true").mode is JobMode.PTY + + with pytest.raises(ValidationError): + RunJobRequest(script="true", mode="stdout") # pyright: ignore[reportArgumentType] + + def test_shell_tool_timeout_budget_has_one_source_of_truth() -> None: assert MAX_WAIT_TIMEOUT_SECONDS == SHELL_TOOL_HARD_TIMEOUT_SECONDS == 300 assert SHELL_TOOL_HTTP_TIMEOUT_GRACE_SECONDS == 10 diff --git a/docker/envs/core-services/api.env.example b/docker/envs/core-services/api.env.example index 9a68cbff70e..e71b4072148 100644 --- a/docker/envs/core-services/api.env.example +++ b/docker/envs/core-services/api.env.example @@ -11,6 +11,8 @@ PLUGIN_REMOTE_INSTALL_PORT=5003 PLUGIN_MAX_PACKAGE_SIZE=52428800 PLUGIN_DAEMON_TIMEOUT=600.0 INNER_API_KEY_FOR_PLUGIN=QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1 +# Client deadline for converting a Binding file to a ToolFile through the Agent backend. +AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS=240 KNOWLEDGE_FS_ENABLED=${KNOWLEDGE_FS_ENABLED:-false} # Production deployments require HTTPS; plain HTTP is limited to non-production or loopback. KNOWLEDGE_FS_BASE_URL= diff --git a/docker/envs/core-services/dify-agent.env.example b/docker/envs/core-services/dify-agent.env.example index 44d5f4ebb0f..ade21650e2a 100644 --- a/docker/envs/core-services/dify-agent.env.example +++ b/docker/envs/core-services/dify-agent.env.example @@ -13,6 +13,7 @@ DIFY_AGENT_SHUTDOWN_GRACE_SECONDS=30 DIFY_AGENT_RUN_RETENTION_SECONDS=259200 # Pydantic AI run deadline; its default matches the independently configurable E2B active timeout. DIFY_AGENT_RUN_TIMEOUT_SECONDS=3600 +DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS=210 # Leave empty to derive from PLUGIN_DAEMON_URL and PLUGIN_DAEMON_KEY in Docker Compose. DIFY_AGENT_PLUGIN_DAEMON_URL= diff --git a/eslint.config.mjs b/eslint.config.mjs index 0bd17450e00..b920fe54030 100644 --- a/eslint.config.mjs +++ b/eslint.config.mjs @@ -2,7 +2,6 @@ import markdown from '@eslint/markdown' import md from 'eslint-markdown' -import hyoban from 'eslint-plugin-hyoban' import jsonc from 'eslint-plugin-jsonc' import markdownPreferences from 'eslint-plugin-markdown-preferences' import pnpm from 'eslint-plugin-pnpm' @@ -421,12 +420,11 @@ export default defineConfig([ files: ['web/i18n/**/*.json'], plugins: { dify, - hyoban, }, rules: { 'dify/consistent-placeholders': 'error', + 'dify/i18n-flat-key': 'error', 'dify/no-extra-keys': 'error', - 'hyoban/i18n-flat-key': 'error', 'jsonc/sort-keys': 'error', }, }, diff --git a/knip.config.ts b/knip.config.ts index 46d5ecc3573..2914f398f8c 100644 --- a/knip.config.ts +++ b/knip.config.ts @@ -15,8 +15,6 @@ const config: KnipConfig = { 'tsslint.config.ts', 'dev-proxy.config.ts', 'plugins/eslint/index.js', - 'vitest.browser.config.ts', - 'vitest.browser.setup.ts', ], project: [ '**/*.{js,mjs,cjs,jsx,ts,tsx,mts,cts,css,mdx}!', @@ -26,6 +24,7 @@ const config: KnipConfig = { '!.storybook/**!', '!plugins/**!', '!test/**!', + '!vitest.browser.setup.ts!', '!vitest.setup.ts!', ], ignore: ['public/**'], diff --git a/lint.config.ts b/lint.config.ts index 29b34bdba03..a5239e811f5 100644 --- a/lint.config.ts +++ b/lint.config.ts @@ -165,6 +165,10 @@ export const lintConfig = { 'eslint-plugin-antfu', ...(enableTailwindCanonicalClasses ? ['eslint-plugin-better-tailwindcss'] : []), 'eslint-plugin-command', + { + name: 'dify', + specifier: './web/plugins/eslint/index.js', + }, 'eslint-plugin-erasable-syntax-only', { name: 'eslint-comments', @@ -174,7 +178,6 @@ export const lintConfig = { name: 'eslint-react', specifier: '@eslint-react/eslint-plugin', }, - 'eslint-plugin-hyoban', { name: 'jsdoc-js', specifier: 'eslint-plugin-jsdoc', @@ -814,7 +817,7 @@ export const lintConfig = { { files: ['web/**/*.tsx'], rules: { - 'hyoban/prefer-tailwind-icons': [ + 'dify/prefer-tailwind-icons': [ 'warn', { prefix: 'i-', diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index 9987dbd7a24..ccd8695c07e 100644 --- a/oxlint-suppressions.json +++ b/oxlint-suppressions.json @@ -699,11 +699,6 @@ "count": 2 } }, - "web/app/components/base/chat/chat-with-history/sidebar/rename-modal.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/app/components/base/chat/chat/__tests__/hooks.spec.tsx": { "no-restricted-imports": { "count": 1 @@ -1922,11 +1917,6 @@ "count": 3 } }, - "web/app/components/datasets/create/empty-dataset-creation-modal/index.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/app/components/datasets/create/file-preview/index.tsx": { "eslint-react/set-state-in-effect": { "count": 1 @@ -2085,11 +2075,6 @@ "count": 1 } }, - "web/app/components/datasets/documents/components/rename-modal.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/app/components/datasets/documents/create-from-pipeline/data-source-options/option-card.tsx": { "jsx-a11y/click-events-have-key-events": { "count": 1 @@ -3362,14 +3347,6 @@ "count": 1 } }, - "web/app/components/workflow/header/online-users.tsx": { - "jsx-a11y/click-events-have-key-events": { - "count": 2 - }, - "jsx-a11y/no-static-element-interactions": { - "count": 2 - } - }, "web/app/components/workflow/header/test-run-menu.tsx": { "erasable-syntax-only/enums": { "count": 1 @@ -5136,11 +5113,6 @@ "count": 7 } }, - "web/app/install/installForm.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/app/reset-password/check-code/page.tsx": { "no-restricted-imports": { "count": 1 @@ -5171,21 +5143,11 @@ "count": 1 } }, - "web/app/signup/components/input-mail.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/app/signup/layout.tsx": { "typescript/no-explicit-any": { "count": 1 } }, - "web/app/signup/set-password/page.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/context/hooks/use-trigger-events-limit-modal.ts": { "eslint-react/set-state-in-effect": { "count": 3 diff --git a/package.json b/package.json index b92b6418e74..c3d019964f3 100644 --- a/package.json +++ b/package.json @@ -37,7 +37,6 @@ "eslint-plugin-better-tailwindcss": "catalog:", "eslint-plugin-command": "catalog:", "eslint-plugin-erasable-syntax-only": "catalog:", - "eslint-plugin-hyoban": "catalog:", "eslint-plugin-jsdoc": "catalog:", "eslint-plugin-jsonc": "catalog:", "eslint-plugin-markdown-preferences": "catalog:", diff --git a/packages/contracts/account-profile-zod.test.ts b/packages/contracts/account-profile-zod.test.ts new file mode 100644 index 00000000000..f55eede9416 --- /dev/null +++ b/packages/contracts/account-profile-zod.test.ts @@ -0,0 +1,13 @@ +import { describe, expect, it } from 'vitest' +import { zAccountProfilePatchPayload } from './generated/api/console/account/zod.gen' + +describe('generated account profile schema', () => { + it('matches the server rules for partial updates', () => { + expect(zAccountProfilePatchPayload.safeParse({ name: 'Jane' }).success).toBe(true) + expect(zAccountProfilePatchPayload.safeParse({}).success).toBe(true) + expect(zAccountProfilePatchPayload.safeParse({ name: null }).success).toBe(false) + expect( + zAccountProfilePatchPayload.safeParse({ name: 'Jane', unexpected: 'value' }).success, + ).toBe(false) + }) +}) diff --git a/packages/contracts/generated/api/console/account/orpc.gen.ts b/packages/contracts/generated/api/console/account/orpc.gen.ts index 74d541a4064..8648b2d6d9a 100644 --- a/packages/contracts/generated/api/console/account/orpc.gen.ts +++ b/packages/contracts/generated/api/console/account/orpc.gen.ts @@ -12,6 +12,8 @@ import { zGetAccountEducationVerifyResponse, zGetAccountIntegratesResponse, zGetAccountProfileResponse, + zPatchAccountProfileBody, + zPatchAccountProfileResponse, zPostAccountAvatarBody, zPostAccountAvatarResponse, zPostAccountChangeEmailBody, @@ -57,8 +59,15 @@ export const get = oc .input(z.object({ query: zGetAccountAvatarQuery })) .output(zGetAccountAvatarResponse) +/** + * Deprecated. Use PATCH /account/profile instead. + * + * @deprecated + */ export const post = oc .route({ + deprecated: true, + description: 'Deprecated. Use PATCH /account/profile instead.', inputStructure: 'detailed', method: 'POST', operationId: 'postAccountAvatar', @@ -268,8 +277,15 @@ export const integrates = { get: get6, } +/** + * Deprecated. Use PATCH /account/profile instead. + * + * @deprecated + */ export const post10 = oc .route({ + deprecated: true, + description: 'Deprecated. Use PATCH /account/profile instead.', inputStructure: 'detailed', method: 'POST', operationId: 'postAccountInterfaceLanguage', @@ -283,8 +299,15 @@ export const interfaceLanguage = { post: post10, } +/** + * Deprecated. Use PATCH /account/profile instead. + * + * @deprecated + */ export const post11 = oc .route({ + deprecated: true, + description: 'Deprecated. Use PATCH /account/profile instead.', inputStructure: 'detailed', method: 'POST', operationId: 'postAccountInterfaceTheme', @@ -298,8 +321,15 @@ export const interfaceTheme = { post: post11, } +/** + * Deprecated. Use PATCH /account/profile instead. + * + * @deprecated + */ export const post12 = oc .route({ + deprecated: true, + description: 'Deprecated. Use PATCH /account/profile instead.', inputStructure: 'detailed', method: 'POST', operationId: 'postAccountName', @@ -338,12 +368,31 @@ export const get7 = oc }) .output(zGetAccountProfileResponse) +export const patch = oc + .route({ + inputStructure: 'detailed', + method: 'PATCH', + operationId: 'patchAccountProfile', + path: '/account/profile', + tags: ['console'], + }) + .input(z.object({ body: zPatchAccountProfileBody })) + .output(zPatchAccountProfileResponse) + export const profile = { get: get7, + patch, } +/** + * Deprecated. Use PATCH /account/profile instead. + * + * @deprecated + */ export const post14 = oc .route({ + deprecated: true, + description: 'Deprecated. Use PATCH /account/profile instead.', inputStructure: 'detailed', method: 'POST', operationId: 'postAccountTimezone', diff --git a/packages/contracts/generated/api/console/account/types.gen.ts b/packages/contracts/generated/api/console/account/types.gen.ts index 882050e4647..a64bd639059 100644 --- a/packages/contracts/generated/api/console/account/types.gen.ts +++ b/packages/contracts/generated/api/console/account/types.gen.ts @@ -125,6 +125,14 @@ export type AccountPasswordPayload = { repeat_new_password: string } +export type AccountProfilePatchPayload = { + avatar?: string + interface_language?: string + interface_theme?: 'dark' | 'light' + name?: string + timezone?: string +} + export type AccountTimezonePayload = { timezone: string } @@ -432,6 +440,20 @@ export type GetAccountProfileResponses = { export type GetAccountProfileResponse = GetAccountProfileResponses[keyof GetAccountProfileResponses] +export type PatchAccountProfileData = { + body: AccountProfilePatchPayload + path?: never + query?: never + url: '/account/profile' +} + +export type PatchAccountProfileResponses = { + 200: AccountResponse +} + +export type PatchAccountProfileResponse = + PatchAccountProfileResponses[keyof PatchAccountProfileResponses] + export type PostAccountTimezoneData = { body: AccountTimezonePayload path?: never diff --git a/packages/contracts/generated/api/console/account/zod.gen.ts b/packages/contracts/generated/api/console/account/zod.gen.ts index 6805df133b2..41e03f24131 100644 --- a/packages/contracts/generated/api/console/account/zod.gen.ts +++ b/packages/contracts/generated/api/console/account/zod.gen.ts @@ -182,6 +182,19 @@ export const zAccountPasswordPayload = z.object({ repeat_new_password: z.string(), }) +/** + * AccountProfilePatchPayload + */ +export const zAccountProfilePatchPayload = z + .object({ + avatar: z.string().optional(), + interface_language: z.string().optional(), + interface_theme: z.enum(['dark', 'light']).optional(), + name: z.string().min(3).max(30).optional(), + timezone: z.string().optional(), + }) + .strict() + /** * AccountTimezonePayload */ @@ -359,6 +372,13 @@ export const zPostAccountPasswordResponse = zAccountResponse */ export const zGetAccountProfileResponse = zAccountResponse +export const zPatchAccountProfileBody = zAccountProfilePatchPayload + +/** + * Success + */ +export const zPatchAccountProfileResponse = zAccountResponse + export const zPostAccountTimezoneBody = zAccountTimezonePayload /** diff --git a/packages/contracts/generated/api/console/explore/types.gen.ts b/packages/contracts/generated/api/console/explore/types.gen.ts index 55181dbc577..965221af60d 100644 --- a/packages/contracts/generated/api/console/explore/types.gen.ts +++ b/packages/contracts/generated/api/console/explore/types.gen.ts @@ -13,7 +13,15 @@ export type LearnDifyAppListResponse = { recommended_apps: Array } -export type RecommendedAppDetailNullableResponse = RecommendedAppDetailResponse | null +export type RecommendedAppDetailResponse = { + can_trial: boolean + export_data: string + icon?: string | null + icon_background?: string | null + id: string + mode: string + name: string +} export type BannerListResponse = Array @@ -30,16 +38,6 @@ export type RecommendedAppResponse = { privacy_policy?: string | null } -export type RecommendedAppDetailResponse = { - can_trial: boolean - export_data: string - icon?: string | null - icon_background?: string | null - id: string - mode: string - name: string -} - export type BannerResponse = { content: BannerContentResponse created_at: string @@ -139,8 +137,12 @@ export type GetExploreAppsByAppIdData = { url: '/explore/apps/{app_id}' } +export type GetExploreAppsByAppIdErrors = { + 404: unknown +} + export type GetExploreAppsByAppIdResponses = { - 200: RecommendedAppDetailNullableResponse + 200: RecommendedAppDetailResponse } export type GetExploreAppsByAppIdResponse = diff --git a/packages/contracts/generated/api/console/explore/zod.gen.ts b/packages/contracts/generated/api/console/explore/zod.gen.ts index efd87b24f91..272b2884146 100644 --- a/packages/contracts/generated/api/console/explore/zod.gen.ts +++ b/packages/contracts/generated/api/console/explore/zod.gen.ts @@ -15,11 +15,6 @@ export const zRecommendedAppDetailResponse = z.object({ name: z.string(), }) -/** - * RecommendedAppDetailNullableResponse - */ -export const zRecommendedAppDetailNullableResponse = zRecommendedAppDetailResponse.nullable() - /** * RecommendedAppInfoResponse */ @@ -166,7 +161,7 @@ export const zGetExploreAppsByAppIdPath = z.object({ /** * Success */ -export const zGetExploreAppsByAppIdResponse = zRecommendedAppDetailNullableResponse +export const zGetExploreAppsByAppIdResponse = zRecommendedAppDetailResponse export const zGetExploreBannersQuery = z.object({ language: z.string().optional().default('en-US'), diff --git a/packages/contracts/generated/api/console/snippets/orpc.gen.ts b/packages/contracts/generated/api/console/snippets/orpc.gen.ts index 433d4259def..f4b77908296 100644 --- a/packages/contracts/generated/api/console/snippets/orpc.gen.ts +++ b/packages/contracts/generated/api/console/snippets/orpc.gen.ts @@ -3,6 +3,8 @@ import { oc } from '@orpc/contract' import * as z from 'zod' import { + zDeleteSnippetsBySnippetIdWorkflowsByWorkflowIdPath, + zDeleteSnippetsBySnippetIdWorkflowsByWorkflowIdResponse, zDeleteSnippetsBySnippetIdWorkflowsDraftNodesByNodeIdVariablesPath, zDeleteSnippetsBySnippetIdWorkflowsDraftNodesByNodeIdVariablesResponse, zDeleteSnippetsBySnippetIdWorkflowsDraftVariablesByVariableIdPath, @@ -874,6 +876,25 @@ export const restore = { post: post12, } +/** + * Delete a published snippet workflow version + * + * Delete a published snippet workflow version + */ +export const delete4 = oc + .route({ + description: 'Delete a published snippet workflow version', + inputStructure: 'detailed', + method: 'DELETE', + operationId: 'deleteSnippetsBySnippetIdWorkflowsByWorkflowId', + path: '/snippets/{snippet_id}/workflows/{workflow_id}', + successStatus: 204, + summary: 'Delete a published snippet workflow version', + tags: ['console'], + }) + .input(z.object({ params: zDeleteSnippetsBySnippetIdWorkflowsByWorkflowIdPath })) + .output(zDeleteSnippetsBySnippetIdWorkflowsByWorkflowIdResponse) + /** * Update a published snippet workflow version's display metadata * @@ -898,6 +919,7 @@ export const patch2 = oc .output(zPatchSnippetsBySnippetIdWorkflowsByWorkflowIdResponse) export const byWorkflowId = { + delete: delete4, patch: patch2, restore, } diff --git a/packages/contracts/generated/api/console/snippets/types.gen.ts b/packages/contracts/generated/api/console/snippets/types.gen.ts index 5dc2a624c1e..da4479f95e4 100644 --- a/packages/contracts/generated/api/console/snippets/types.gen.ts +++ b/packages/contracts/generated/api/console/snippets/types.gen.ts @@ -1758,6 +1758,28 @@ export type PostSnippetsBySnippetIdWorkflowsPublishResponses = { export type PostSnippetsBySnippetIdWorkflowsPublishResponse = PostSnippetsBySnippetIdWorkflowsPublishResponses[keyof PostSnippetsBySnippetIdWorkflowsPublishResponses] +export type DeleteSnippetsBySnippetIdWorkflowsByWorkflowIdData = { + body?: never + path: { + snippet_id: string + workflow_id: string + } + query?: never + url: '/snippets/{snippet_id}/workflows/{workflow_id}' +} + +export type DeleteSnippetsBySnippetIdWorkflowsByWorkflowIdErrors = { + 400: unknown + 404: unknown +} + +export type DeleteSnippetsBySnippetIdWorkflowsByWorkflowIdResponses = { + 204: void +} + +export type DeleteSnippetsBySnippetIdWorkflowsByWorkflowIdResponse = + DeleteSnippetsBySnippetIdWorkflowsByWorkflowIdResponses[keyof DeleteSnippetsBySnippetIdWorkflowsByWorkflowIdResponses] + export type PatchSnippetsBySnippetIdWorkflowsByWorkflowIdData = { body: WorkflowUpdatePayload path: { diff --git a/packages/contracts/generated/api/console/snippets/zod.gen.ts b/packages/contracts/generated/api/console/snippets/zod.gen.ts index 994adf7aa0f..65ac302b2ae 100644 --- a/packages/contracts/generated/api/console/snippets/zod.gen.ts +++ b/packages/contracts/generated/api/console/snippets/zod.gen.ts @@ -1988,6 +1988,16 @@ export const zPostSnippetsBySnippetIdWorkflowsPublishPath = z.object({ */ export const zPostSnippetsBySnippetIdWorkflowsPublishResponse = zWorkflowPublishResponse +export const zDeleteSnippetsBySnippetIdWorkflowsByWorkflowIdPath = z.object({ + snippet_id: z.uuid(), + workflow_id: z.string(), +}) + +/** + * Workflow deleted successfully + */ +export const zDeleteSnippetsBySnippetIdWorkflowsByWorkflowIdResponse = z.void() + export const zPatchSnippetsBySnippetIdWorkflowsByWorkflowIdBody = zWorkflowUpdatePayload export const zPatchSnippetsBySnippetIdWorkflowsByWorkflowIdPath = z.object({ diff --git a/packages/contracts/generated/api/console/tags/orpc.gen.ts b/packages/contracts/generated/api/console/tags/orpc.gen.ts index f315d1272e5..2398d6517b4 100644 --- a/packages/contracts/generated/api/console/tags/orpc.gen.ts +++ b/packages/contracts/generated/api/console/tags/orpc.gen.ts @@ -50,7 +50,7 @@ export const get = oc path: '/tags', tags: ['console'], }) - .input(z.object({ query: zGetTagsQuery.optional() })) + .input(z.object({ query: zGetTagsQuery })) .output(zGetTagsResponse) export const post = oc diff --git a/packages/contracts/generated/api/console/tags/types.gen.ts b/packages/contracts/generated/api/console/tags/types.gen.ts index 14c7c9c722b..4bc2ce8e949 100644 --- a/packages/contracts/generated/api/console/tags/types.gen.ts +++ b/packages/contracts/generated/api/console/tags/types.gen.ts @@ -27,9 +27,9 @@ export type TagType = 'app' | 'knowledge' | 'snippet' export type GetTagsData = { body?: never path?: never - query?: { + query: { keyword?: string - type?: '' | 'app' | 'knowledge' | 'snippet' + type: 'app' | 'knowledge' | 'snippet' } url: '/tags' } diff --git a/packages/contracts/generated/api/console/tags/zod.gen.ts b/packages/contracts/generated/api/console/tags/zod.gen.ts index 7cab6b4df9e..13d8259cc58 100644 --- a/packages/contracts/generated/api/console/tags/zod.gen.ts +++ b/packages/contracts/generated/api/console/tags/zod.gen.ts @@ -41,7 +41,7 @@ export const zTagBasePayload = z.object({ export const zGetTagsQuery = z.object({ keyword: z.string().optional(), - type: z.enum(['', 'app', 'knowledge', 'snippet']).optional().default(''), + type: z.enum(['app', 'knowledge', 'snippet']), }) /** diff --git a/packages/contracts/generated/api/console/workspaces/types.gen.ts b/packages/contracts/generated/api/console/workspaces/types.gen.ts index a7c53ab7a4b..99dab0f6291 100644 --- a/packages/contracts/generated/api/console/workspaces/types.gen.ts +++ b/packages/contracts/generated/api/console/workspaces/types.gen.ts @@ -2699,6 +2699,7 @@ export type GetWorkspacesCurrentCustomizedSnippetsBySnippetIdExportData = { } query?: { include_secret?: string + workflow_id?: string } url: '/workspaces/current/customized-snippets/{snippet_id}/export' } diff --git a/packages/contracts/generated/api/console/workspaces/zod.gen.ts b/packages/contracts/generated/api/console/workspaces/zod.gen.ts index 8bcd255b3c9..c55c0e70099 100644 --- a/packages/contracts/generated/api/console/workspaces/zod.gen.ts +++ b/packages/contracts/generated/api/console/workspaces/zod.gen.ts @@ -3609,6 +3609,7 @@ export const zGetWorkspacesCurrentCustomizedSnippetsBySnippetIdExportPath = z.ob export const zGetWorkspacesCurrentCustomizedSnippetsBySnippetIdExportQuery = z.object({ include_secret: z.string().optional().default('false'), + workflow_id: z.string().optional(), }) /** diff --git a/packages/contracts/generated/api/web/types.gen.ts b/packages/contracts/generated/api/web/types.gen.ts index e73d1e31fcb..7675ff98af2 100644 --- a/packages/contracts/generated/api/web/types.gen.ts +++ b/packages/contracts/generated/api/web/types.gen.ts @@ -1569,6 +1569,7 @@ export type GetWebappPermissionErrors = { 400: unknown 401: unknown 500: unknown + 503: unknown } export type GetWebappPermissionResponses = { diff --git a/packages/contracts/openapi-ts.api.config.ts b/packages/contracts/openapi-ts.api.config.ts index f77057b1028..d34cdd99080 100644 --- a/packages/contracts/openapi-ts.api.config.ts +++ b/packages/contracts/openapi-ts.api.config.ts @@ -61,6 +61,7 @@ const currentDir = path.dirname(fileURLToPath(import.meta.url)) const apiOpenApiDir = path.resolve(currentDir, 'openapi') const operationMethods = new Set(['delete', 'get', 'patch', 'post', 'put']) +const strictZodSchemaNames = new Set(['AccountProfilePatchPayload']) const pydanticDecimalStringPattern = '^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$' const codegenSafeDecimalStringPattern = '^(?![-+.]*$)[+-]?0*\\d*\\.?\\d*$' const fastOpenApiConsoleSpecFilename = 'fastopenapi-console-openapi.json' @@ -492,6 +493,22 @@ const createApiConfig = (job: ApiJob): UserConfig => ({ { name: 'zod', '~resolvers': { + object: (ctx) => { + const objectSchema = ctx.nodes.base(ctx) + const additionalProperties = ctx.schema.additionalProperties + // openapi-ts normalizes `additionalProperties: false` to `never`, but + // does not make shaped Zod objects strict. + const isStrictSchema = ctx.path['~ref'].some( + (segment) => typeof segment === 'string' && strictZodSchemaNames.has(segment), + ) + if ( + isStrictSchema && + (additionalProperties === false || additionalProperties?.type === 'never') + ) + return objectSchema.attr('strict').call() + + return objectSchema + }, string: (ctx) => { if (ctx.schema.format === 'binary') return $(ctx.symbols.z) diff --git a/packages/dify-ui/.gitignore b/packages/dify-ui/.gitignore index befe881885d..b4b3fdf70f9 100644 --- a/packages/dify-ui/.gitignore +++ b/packages/dify-ui/.gitignore @@ -1,3 +1,4 @@ /coverage +/.vitest-browser /dist /storybook-static diff --git a/packages/dify-ui/README.md b/packages/dify-ui/README.md index 75c35061387..1f0987dfc40 100644 --- a/packages/dify-ui/README.md +++ b/packages/dify-ui/README.md @@ -276,11 +276,13 @@ See the [web overlay guide] for the web app overlay best practices. ## Development -- `vp check packages/dify-ui` (from the repository root) — formatting and lint for the package plus the repository-wide TypeScript diagnostics configured by Vite+. -- `pnpm -C packages/dify-ui test` — Vitest unit tests for primitives. -- `pnpm -C packages/dify-ui storybook` — Storybook on the default port. Each primitive has `index.stories.tsx`. -- `pnpm -C packages/dify-ui test:storybook` — Storybook component tests in Vitest browser mode. Stories without `play` are render and a11y smoke tests; stories with `play` should cover public UI contracts such as opening overlays, keyboard navigation, disabled/loading guards, form submission, and controlled state updates. -- `pnpm -C packages/dify-ui type-check` — TypeScript 7 native type checking for this package only. +Run `vp check packages/dify-ui` from the repository root for package formatting, lint, and repository-wide TypeScript diagnostics. Run the remaining commands from `packages/dify-ui/`: + +- `vp test --project unit` — Vitest unit tests for primitives. +- `vp run storybook` — Storybook on the default port. Each primitive has `index.stories.tsx`. +- `vp test --project storybook --run` — Storybook component tests in Vitest browser mode. Stories without `play` are render and a11y smoke tests; stories with `play` should cover public UI contracts such as opening overlays, keyboard navigation, disabled/loading guards, form submission, and controlled state updates. + +Both test projects run in Playwright Chromium Browser Mode; choose `unit` or `storybook` by behavior owner, not runtime. Bare `vp test` runs both projects. ### Test Boundary @@ -291,10 +293,10 @@ wrapper contracts such as class variants, Base UI passthrough props, hidden inpu serialization, data attribute hooks, store behavior, and edge cases that do not need a full story. -Storybook accessibility testing stays enabled globally with `a11y.test = 'error'`. -If a story is temporarily marked `todo`, keep the exception local to that story -and do not treat an interaction `play` test as a replacement for fixing the -underlying accessibility issue. +Storybook accessibility testing uses `a11y.test = 'error'` for enabled rules. +Color contrast is a known design-token gap and is currently excluded globally; +do not add another global exclusion. Keep other temporary exceptions local to +the affected story and do not use a `play` test in place of an accessibility fix. ### Disabling Animations In Tests @@ -311,6 +313,7 @@ Set the Base UI test flag in a Vitest setup file to skip those waits: ``` `packages/dify-ui/vitest.setup.ts` already applies this for primitive tests. +The Storybook project intentionally uses its preview setup instead; do not disable animation lifecycles globally there. See [component authoring rules] for: diff --git a/packages/dify-ui/src/avatar/__tests__/index.spec.tsx b/packages/dify-ui/src/avatar/__tests__/index.spec.tsx index 8231d196e5d..c7af2a633cd 100644 --- a/packages/dify-ui/src/avatar/__tests__/index.spec.tsx +++ b/packages/dify-ui/src/avatar/__tests__/index.spec.tsx @@ -1,12 +1,24 @@ import { render } from 'vitest-browser-react' import { Avatar } from '..' +const avatarDataUrl = 'data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP///ywAAAAAAQABAAACAUwAOw==' + function stubImageLoader() { const originalImage = window.Image const images: HTMLImageElement[] = [] function TestImage(_width?: number, _height?: number): HTMLImageElement { - const image = document.createElement('img') + const image = { + complete: false, + crossOrigin: null, + naturalWidth: 0, + onerror: null, + onload: null, + referrerPolicy: '', + sizes: '', + src: '', + srcset: '', + } as unknown as HTMLImageElement images.push(image) return image } @@ -51,11 +63,7 @@ describe('Avatar', () => { try { const screen = await render( - , + , ) await expect.element(screen.getByText('J')).toBeVisible() @@ -63,7 +71,7 @@ describe('Avatar', () => { expect(onStatusChange).toHaveBeenCalledWith('loading') }) - images[0]?.onload?.(new Event('load')) + images.at(-1)?.onload?.(new Event('load')) await vi.waitFor(() => { expect(onStatusChange).toHaveBeenCalledWith('loaded') diff --git a/packages/dify-ui/src/context-menu/__tests__/index.spec.tsx b/packages/dify-ui/src/context-menu/__tests__/index.spec.tsx index 4e6c7c79e4e..336c7eb4b55 100644 --- a/packages/dify-ui/src/context-menu/__tests__/index.spec.tsx +++ b/packages/dify-ui/src/context-menu/__tests__/index.spec.tsx @@ -98,7 +98,7 @@ describe('context-menu wrapper', () => { .element(screen.getByRole('group', { name: 'context content positioner' })) .toHaveAttribute('id', 'context-content-positioner') await expect.element(screen.getByRole('menu')).toHaveAttribute('id', 'context-content-popup') - expect(handlePositionerMouseEnter).toHaveBeenCalledTimes(1) + expect(handlePositionerMouseEnter).toHaveBeenCalled() expect(handlePopupClick).toHaveBeenCalledTimes(1) }) }) diff --git a/packages/dify-ui/src/toast/__tests__/index.spec.tsx b/packages/dify-ui/src/toast/__tests__/index.spec.tsx index 36fc4d14f3b..8c690c134dd 100644 --- a/packages/dify-ui/src/toast/__tests__/index.spec.tsx +++ b/packages/dify-ui/src/toast/__tests__/index.spec.tsx @@ -1,5 +1,5 @@ +import { userEvent } from 'vite-plus/test/browser' import { render } from 'vitest-browser-react' -import { userEvent } from 'vitest/browser' import { createToast, createToastManager, toast, ToastHost } from '../index' const asHTMLElement = (element: HTMLElement | SVGElement) => element as HTMLElement diff --git a/packages/dify-ui/tsconfig.json b/packages/dify-ui/tsconfig.json index 3dac6d6f3f5..c9c9e935410 100644 --- a/packages/dify-ui/tsconfig.json +++ b/packages/dify-ui/tsconfig.json @@ -9,7 +9,6 @@ "src/**/*.ts", "src/**/*.tsx", "vite.config.ts", - "vitest.config.ts", "vitest.setup.ts" ], "exclude": ["node_modules", "dist", "storybook-static", "coverage"] diff --git a/packages/dify-ui/vite.config.ts b/packages/dify-ui/vite.config.ts index b5e82761bd3..30bca171649 100644 --- a/packages/dify-ui/vite.config.ts +++ b/packages/dify-ui/vite.config.ts @@ -1,9 +1,71 @@ +import path from 'node:path' +import { fileURLToPath } from 'node:url' +import { storybookTest } from '@storybook/addon-vitest/vitest-plugin' +import tailwindcss from '@tailwindcss/vite' import react from '@vitejs/plugin-react' import { defineConfig } from 'vite-plus' +import { playwright } from 'vite-plus/test/browser-playwright' + +const dirname = path.dirname(fileURLToPath(import.meta.url)) +const configDir = path.join(dirname, '.storybook') +const isCI = !!process.env.CI export default defineConfig({ plugins: [react()], resolve: { tsconfigPaths: true, }, + optimizeDeps: { + include: ['vite-plus/test/browser'], + }, + test: { + browser: { + enabled: true, + provider: playwright(), + instances: [{ browser: 'chromium' }], + headless: true, + screenshotDirectory: './.vitest-browser/screenshots', + screenshotFailures: true, + }, + coverage: { + provider: 'v8', + include: ['src/**/*.{ts,tsx}'], + exclude: [ + 'src/**/*.stories.{ts,tsx}', + 'src/**/__tests__/**', + 'src/themes/**', + 'src/styles/**', + ], + reporter: isCI ? ['json', 'json-summary'] : ['text', 'json', 'json-summary'], + }, + projects: [ + { + extends: true, + plugins: [tailwindcss()], + test: { + name: 'unit', + globals: true, + setupFiles: ['./vitest.setup.ts'], + include: ['src/**/__tests__/**/*.spec.{ts,tsx}'], + browser: { + trace: { + mode: 'retain-on-failure', + tracesDir: './.vitest-browser/traces', + }, + }, + }, + }, + { + extends: true, + plugins: [ + storybookTest({ + configDir, + }), + ], + test: { + name: 'storybook', + }, + }, + ], + }, }) diff --git a/packages/dify-ui/vitest.config.ts b/packages/dify-ui/vitest.config.ts deleted file mode 100644 index 214c8759432..00000000000 --- a/packages/dify-ui/vitest.config.ts +++ /dev/null @@ -1,69 +0,0 @@ -import path from 'node:path' -import { fileURLToPath } from 'node:url' -import { storybookTest } from '@storybook/addon-vitest/vitest-plugin' -import tailwindcss from '@tailwindcss/vite' -import react from '@vitejs/plugin-react' -import { defineConfig } from 'vite-plus' -import { playwright } from 'vite-plus/test/browser-playwright' - -const dirname = path.dirname(fileURLToPath(import.meta.url)) -const configDir = path.join(dirname, '.storybook') -const isCI = !!process.env.CI - -export default defineConfig({ - plugins: [react()], - resolve: { - tsconfigPaths: true, - }, - optimizeDeps: { - include: ['vite-plus/test/browser'], - }, - test: { - coverage: { - provider: 'v8', - include: ['src/**/*.{ts,tsx}'], - exclude: [ - 'src/**/*.stories.{ts,tsx}', - 'src/**/__tests__/**', - 'src/themes/**', - 'src/styles/**', - ], - reporter: isCI ? ['json', 'json-summary'] : ['text', 'json', 'json-summary'], - }, - projects: [ - { - extends: true, - plugins: [tailwindcss()], - test: { - name: 'unit', - globals: true, - setupFiles: ['./vitest.setup.ts'], - include: ['src/**/__tests__/**/*.spec.{ts,tsx}'], - browser: { - enabled: true, - provider: playwright(), - instances: [{ browser: 'chromium' }], - headless: true, - }, - }, - }, - { - extends: true, - plugins: [ - storybookTest({ - configDir, - }), - ], - test: { - name: 'storybook', - browser: { - enabled: true, - provider: playwright(), - instances: [{ browser: 'chromium' }], - headless: true, - }, - }, - }, - ], - }, -}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 51bc79a5093..96ba1f25d8f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -330,9 +330,6 @@ catalogs: eslint-plugin-erasable-syntax-only: specifier: 0.4.2 version: 0.4.2 - eslint-plugin-hyoban: - specifier: 0.14.1 - version: 0.14.1 eslint-plugin-jsdoc: specifier: 63.3.3 version: 63.3.3 @@ -693,9 +690,6 @@ importers: eslint-plugin-erasable-syntax-only: specifier: 'catalog:' version: 0.4.2(@typescript-eslint/parser@8.67.0(@typescript/typescript6@6.0.2)(eslint@10.8.1(jiti@2.7.0)(supports-color@11.0.0))(supports-color@11.0.0))(@typescript/typescript6@6.0.2)(eslint@10.8.1(jiti@2.7.0)(supports-color@11.0.0))(supports-color@10.2.2) - eslint-plugin-hyoban: - specifier: 'catalog:' - version: 0.14.1(eslint@10.8.1(jiti@2.7.0)(supports-color@11.0.0)) eslint-plugin-jsdoc: specifier: 'catalog:' version: 63.3.3(eslint@10.8.1(jiti@2.7.0)(supports-color@11.0.0))(supports-color@11.0.0) @@ -6018,11 +6012,6 @@ packages: peerDependencies: eslint: '>=8' - eslint-plugin-hyoban@0.14.1: - resolution: {integrity: sha512-R7UX1AMUilGfFftGoHKTlG0BVN5PsiZLN78Yqi6GZBaheQkvwRj4Dw+k+wW+1nKcueyh4IKdvt+n+0ayLEnZYA==} - peerDependencies: - eslint: '*' - eslint-plugin-jsdoc@63.3.3: resolution: {integrity: sha512-xI4IeVRzRFA2DGHrPLIxF3U+oJHU3FE+P9Zb27fVs5dPHgfcpoAs0PyCbznVhK7pwR+9BPUztFeXSgpw/CL4Yg==} engines: {node: ^22.13.0 || >=24} @@ -13430,10 +13419,6 @@ snapshots: eslint: 10.8.1(jiti@2.7.0)(supports-color@11.0.0) eslint-compat-utils: 0.5.1(eslint@10.8.1(jiti@2.7.0)(supports-color@11.0.0)) - eslint-plugin-hyoban@0.14.1(eslint@10.8.1(jiti@2.7.0)(supports-color@11.0.0)): - dependencies: - eslint: 10.8.1(jiti@2.7.0)(supports-color@11.0.0) - eslint-plugin-jsdoc@63.3.3(eslint@10.8.1(jiti@2.7.0)(supports-color@11.0.0))(supports-color@11.0.0): dependencies: '@es-joy/jsdoccomment': 0.91.0 @@ -17186,7 +17171,6 @@ time: eslint-plugin-better-tailwindcss@4.7.0: '2026-07-19T13:26:01.366Z' eslint-plugin-command@3.5.3: '2026-07-14T04:59:31.991Z' eslint-plugin-erasable-syntax-only@0.4.2: '2026-06-18T00:56:01.038Z' - eslint-plugin-hyoban@0.14.1: '2026-03-08T02:51:00.805Z' eslint-plugin-jsdoc@63.3.3: '2026-08-02T18:02:10.488Z' eslint-plugin-jsonc@3.4.1: '2026-08-04T07:21:02.632Z' eslint-plugin-markdown-preferences@0.41.1: '2026-04-09T23:28:41.552Z' diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 091fea64959..17d8f8c5978 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -171,7 +171,6 @@ catalog: eslint-plugin-better-tailwindcss: 4.7.0 eslint-plugin-command: 3.5.3 eslint-plugin-erasable-syntax-only: 0.4.2 - eslint-plugin-hyoban: 0.14.1 eslint-plugin-jsdoc: 63.3.3 eslint-plugin-jsonc: 3.4.1 eslint-plugin-markdown-preferences: 0.41.1 diff --git a/web/README.md b/web/README.md index 3a268f512cb..cd4cc161b1e 100644 --- a/web/README.md +++ b/web/README.md @@ -99,20 +99,22 @@ Then follow the [Lint Documentation] to lint the code. We use [Vitest] and [React Testing Library] for Unit Testing. -**📖 Frontend Testing Guide**: See [web/docs/test.md] for the canonical testing policy and workflow. +**📖 Frontend Testing Guide**: See the [Frontend Testing Guide] for the canonical testing policy and workflow. > [!IMPORTANT] > As we are using Vite+, the `vitest` command is not available. > Please make sure to run tests with `vp` commands. -> For example, use `npx vp test` instead of `npx vitest`. +> For example, use `vp test` instead of `vitest`. Run test: ```bash cd web -vp test run +vp test run --project unit ``` +The standard unit command runs in `happy-dom`. Browser Mode is reserved for behavior that depends on a real browser; see the [Frontend Testing Guide] for its admission criteria and commands. Always select a project explicitly: bare `vp test` runs every registered project, including Browser Mode. + If a test fails only in CI, inspect the failing job and reproduce it locally when possible. A rerun can help identify a flaky test, but it does not replace diagnosing or reporting the failure. ## Documentation @@ -124,6 +126,7 @@ Visit to view the full documentation. The Dify community can be found on [Discord community], where you can ask questions, voice ideas, and share your projects. [Discord community]: https://discord.gg/5AEfbxcd9k +[Frontend Testing Guide]: ./docs/test.md [Lint Documentation]: ./docs/lint.md [Next.js]: https://nextjs.org [Node.js]: https://nodejs.org @@ -133,4 +136,3 @@ The Dify community can be found on [Discord community], where you can ask questi [Vitest]: https://vitest.dev [pnpm]: https://pnpm.io [vinext]: https://github.com/cloudflare/vinext -[web/docs/test.md]: ./docs/test.md diff --git a/web/__tests__/billing/pricing-modal-flow.test.tsx b/web/__tests__/billing/pricing-modal-flow.test.tsx index 1ae301976b9..7bbea49505f 100644 --- a/web/__tests__/billing/pricing-modal-flow.test.tsx +++ b/web/__tests__/billing/pricing-modal-flow.test.tsx @@ -189,12 +189,10 @@ describe('Pricing Modal Flow', () => { expect(screen.getByText(/plansCommon\.annualBilling/i)).toBeInTheDocument() }) - it('should show tax tip in footer for cloud category', () => { + it('should show the tax exclusion notice in the footer for cloud category', () => { render() - // Use exact match to avoid matching taxTipSecond expect(screen.getByText('billing.plansCommon.taxTip')).toBeInTheDocument() - expect(screen.getByText('billing.plansCommon.taxTipSecond')).toBeInTheDocument() }) }) diff --git a/web/app/components/base/chat/chat-with-history/sidebar/__tests__/rename-modal.spec.tsx b/web/app/components/base/chat/chat-with-history/sidebar/__tests__/rename-modal.spec.tsx index aff1a1bc7bc..c3a135bd424 100644 --- a/web/app/components/base/chat/chat-with-history/sidebar/__tests__/rename-modal.spec.tsx +++ b/web/app/components/base/chat/chat-with-history/sidebar/__tests__/rename-modal.spec.tsx @@ -22,8 +22,7 @@ describe('RenameModal', () => { render() expect(screen.getByText('common.chat.renameConversation')).toBeInTheDocument() - expect(screen.getByText('common.chat.conversationName')).toBeInTheDocument() - expect(screen.getByPlaceholderText('common.chat.conversationNamePlaceholder')).toHaveValue( + expect(screen.getByRole('textbox', { name: 'common.chat.conversationName' })).toHaveValue( 'Original Name', ) expect(screen.getByText('common.operation.cancel')).toBeInTheDocument() @@ -50,11 +49,21 @@ describe('RenameModal', () => { const input = screen.getByRole('textbox') await user.clear(input) await user.type(input, 'Updated Name') - await user.click(screen.getByText('common.operation.save')) + await user.keyboard('{Enter}') expect(defaultProps.onSave).toHaveBeenCalledWith('Updated Name') }) + it('does not resubmit while save is pending', async () => { + const user = userEvent.setup() + render() + + await user.click(screen.getByRole('textbox', { name: 'common.chat.conversationName' })) + await user.keyboard('{Enter}') + + expect(defaultProps.onSave).not.toHaveBeenCalled() + }) + it('calls onSave with initial name when unchanged', async () => { const user = userEvent.setup() render() diff --git a/web/app/components/base/chat/chat-with-history/sidebar/rename-modal.tsx b/web/app/components/base/chat/chat-with-history/sidebar/rename-modal.tsx index 85fd0b61467..d9bdfb027bc 100644 --- a/web/app/components/base/chat/chat-with-history/sidebar/rename-modal.tsx +++ b/web/app/components/base/chat/chat-with-history/sidebar/rename-modal.tsx @@ -2,10 +2,12 @@ import type { FC } from 'react' import { Button } from '@langgenius/dify-ui/button' import { Dialog, DialogContent, DialogTitle } from '@langgenius/dify-ui/dialog' +import { Field, FieldLabel } from '@langgenius/dify-ui/field' +import { Form } from '@langgenius/dify-ui/form' +import { Input } from '@langgenius/dify-ui/input' import * as React from 'react' import { useState } from 'react' import { useTranslation } from 'react-i18next' -import Input from '@/app/components/base/input' type IRenameModalProps = { isShow: boolean @@ -27,29 +29,32 @@ const RenameModal: FC = ({ isShow, saveLoading, name, onClose {t(($) => $['chat.renameConversation'], { ns: 'common' })} -
- {t(($) => $['chat.conversationName'], { ns: 'common' })} -
- setTempName(e.target.value)} - placeholder={conversationNamePlaceholder} - /> +
{ + if (!saveLoading) onSave(tempName) + }} + > + + + {t(($) => $['chat.conversationName'], { ns: 'common' })} + + + -
- - -
+
+ + +
+
) diff --git a/web/app/components/base/file-uploader/file-uploader-in-chat-input/__tests__/file-image-item.browser.spec.tsx b/web/app/components/base/file-uploader/file-uploader-in-chat-input/__tests__/file-image-item.browser.spec.tsx index 7e7d4f8e88b..a3f4ab0dea9 100644 --- a/web/app/components/base/file-uploader/file-uploader-in-chat-input/__tests__/file-image-item.browser.spec.tsx +++ b/web/app/components/base/file-uploader/file-uploader-in-chat-input/__tests__/file-image-item.browser.spec.tsx @@ -24,8 +24,10 @@ describe('FileImageItem pointer interaction', () => { it('keeps the preview clickable when the download action is visible', async () => { const screen = await render() const preview = screen.getByRole('button', { name: 'common.operation.view photo.png' }) + const download = screen.getByRole('button', { name: 'common.operation.download' }) await preview.hover() + await expect.element(download).toBeVisible() await preview.click() await expect.element(page.getByRole('dialog')).toBeVisible() diff --git a/web/app/components/billing/pricing/footer.tsx b/web/app/components/billing/pricing/footer.tsx index 9cceb3e6428..bedbd89b664 100644 --- a/web/app/components/billing/pricing/footer.tsx +++ b/web/app/components/billing/pricing/footer.tsx @@ -26,9 +26,6 @@ const Footer = ({ pricingPageURL, currentCategory }: FooterProps) => { {t(($) => $['plansCommon.taxTip'], { ns: 'billing' })} - - {t(($) => $['plansCommon.taxTipSecond'], { ns: 'billing' })} - )} diff --git a/web/app/components/datasets/create/empty-dataset-creation-modal/__tests__/index.spec.tsx b/web/app/components/datasets/create/empty-dataset-creation-modal/__tests__/index.spec.tsx index 912ee8a52d1..884fc5b9446 100644 --- a/web/app/components/datasets/create/empty-dataset-creation-modal/__tests__/index.spec.tsx +++ b/web/app/components/datasets/create/empty-dataset-creation-modal/__tests__/index.spec.tsx @@ -1,5 +1,6 @@ import type { MockedFunction } from 'vite-plus/test' import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import userEvent from '@testing-library/user-event' import { createEmptyDataset } from '@/service/datasets' import { useInvalidDatasetList } from '@/service/knowledge/use-dataset' import EmptyDatasetCreationModal from '../index' @@ -73,9 +74,8 @@ describe('EmptyDatasetCreationModal', () => { expect(screen.getByText('datasetCreation.stepOne.modal.title')).toBeInTheDocument() expect(screen.getByText('datasetCreation.stepOne.modal.tip')).toBeInTheDocument() - expect(screen.getByText('datasetCreation.stepOne.modal.input')).toBeInTheDocument() expect( - screen.getByPlaceholderText('datasetCreation.stepOne.modal.placeholder'), + screen.getByRole('textbox', { name: 'datasetCreation.stepOne.modal.input' }), ).toBeInTheDocument() expect(screen.getByText('datasetCreation.stepOne.modal.confirmButton')).toBeInTheDocument() expect(screen.getByText('datasetCreation.stepOne.modal.cancelButton')).toBeInTheDocument() @@ -289,20 +289,49 @@ describe('EmptyDatasetCreationModal', () => { // API Calls - Test API interactions describe('API Calls', () => { - it('should call createEmptyDataset with correct parameters', async () => { + it('should submit from the dataset name input with Enter', async () => { const mockOnHide = vi.fn() render() - const input = screen.getByPlaceholderText('datasetCreation.stepOne.modal.placeholder') - const confirmButton = screen.getByText('datasetCreation.stepOne.modal.confirmButton') + const user = userEvent.setup() + const input = screen.getByRole('textbox', { name: 'datasetCreation.stepOne.modal.input' }) - fireEvent.change(input, { target: { value: 'New Dataset' } }) - fireEvent.click(confirmButton) + await user.type(input, 'New Dataset{Enter}') await waitFor(() => { expect(mockCreateEmptyDataset).toHaveBeenCalledWith({ name: 'New Dataset' }) }) }) + it('should not submit again while dataset creation is pending', async () => { + let resolveRequest: + | ((value: Awaited>) => void) + | undefined + mockCreateEmptyDataset.mockReturnValueOnce( + new Promise((resolve) => { + resolveRequest = resolve + }), + ) + const onHide = vi.fn() + render() + const user = userEvent.setup() + const input = screen.getByRole('textbox', { name: 'datasetCreation.stepOne.modal.input' }) + + await user.type(input, 'New Dataset{Enter}') + await waitFor(() => { + expect(mockCreateEmptyDataset).toHaveBeenCalledTimes(1) + }) + await user.keyboard('{Enter}') + expect(mockCreateEmptyDataset).toHaveBeenCalledTimes(1) + + resolveRequest?.({ + id: 'dataset-123', + name: 'New Dataset', + } as Awaited>) + await waitFor(() => { + expect(onHide).toHaveBeenCalledTimes(1) + }) + }) + it('should call invalidDatasetList after successful creation', async () => { const mockOnHide = vi.fn() render() diff --git a/web/app/components/datasets/create/empty-dataset-creation-modal/index.module.css b/web/app/components/datasets/create/empty-dataset-creation-modal/index.module.css index c5370ce6507..d284faa2206 100644 --- a/web/app/components/datasets/create/empty-dataset-creation-modal/index.module.css +++ b/web/app/components/datasets/create/empty-dataset-creation-modal/index.module.css @@ -26,7 +26,7 @@ @apply mb-8; } .form .label { - @apply mb-2 text-text-primary; + @apply py-0 text-text-primary; font-weight: 500; font-size: 14px; line-height: 20px; diff --git a/web/app/components/datasets/create/empty-dataset-creation-modal/index.tsx b/web/app/components/datasets/create/empty-dataset-creation-modal/index.tsx index b55c9f3e0c2..c725e3a7d17 100644 --- a/web/app/components/datasets/create/empty-dataset-creation-modal/index.tsx +++ b/web/app/components/datasets/create/empty-dataset-creation-modal/index.tsx @@ -1,13 +1,15 @@ 'use client' import { Button } from '@langgenius/dify-ui/button' import { cn } from '@langgenius/dify-ui/cn' -import { Dialog, DialogContent } from '@langgenius/dify-ui/dialog' +import { Dialog, DialogContent, DialogTitle } from '@langgenius/dify-ui/dialog' +import { Field, FieldLabel } from '@langgenius/dify-ui/field' +import { Form } from '@langgenius/dify-ui/form' +import { Input } from '@langgenius/dify-ui/input' import { toast } from '@langgenius/dify-ui/toast' import * as React from 'react' import { useState } from 'react' import { useTranslation } from 'react-i18next' import { trackEvent } from '@/app/components/base/amplitude' -import Input from '@/app/components/base/input' import { useRouter } from '@/next/navigation' import { createEmptyDataset } from '@/service/datasets' import { useInvalidDatasetList } from '@/service/knowledge/use-dataset' @@ -19,10 +21,13 @@ type IProps = { } const EmptyDatasetCreationModal = ({ show = false, onHide }: IProps) => { const [inputValue, setInputValue] = useState('') + const [isSubmitting, setIsSubmitting] = useState(false) const { t } = useTranslation() const router = useRouter() const invalidDatasetList = useInvalidDatasetList() const submit = async () => { + if (isSubmitting) return + if (!inputValue) { toast.error(t(($) => $['stepOne.modal.nameNotEmpty'], { ns: 'datasetCreation' })) return @@ -31,6 +36,7 @@ const EmptyDatasetCreationModal = ({ show = false, onHide }: IProps) => { toast.error(t(($) => $['stepOne.modal.nameLengthInvalid'], { ns: 'datasetCreation' })) return } + setIsSubmitting(true) try { const dataset = await createEmptyDataset({ name: inputValue }) invalidDatasetList() @@ -42,6 +48,8 @@ const EmptyDatasetCreationModal = ({ show = false, onHide }: IProps) => { router.push(`/datasets/${dataset.id}/documents`) } catch { toast.error(t(($) => $['stepOne.modal.failed'], { ns: 'datasetCreation' })) + } finally { + setIsSubmitting(false) } } return ( @@ -53,9 +61,9 @@ const EmptyDatasetCreationModal = ({ show = false, onHide }: IProps) => { >
-
+ {t(($) => $['stepOne.modal.title'], { ns: 'datasetCreation' })} -
+
{t(($) => $['stepOne.modal.tip'], { ns: 'datasetCreation' })}
-
-
- {t(($) => $['stepOne.modal.input'], { ns: 'datasetCreation' })} +
void submit()}> + + + {t(($) => $['stepOne.modal.input'], { ns: 'datasetCreation' })} + + $['stepOne.modal.placeholder'], { ns: 'datasetCreation' }) || '' + } + onValueChange={setInputValue} + /> + +
+ +
- $['stepOne.modal.placeholder'], { ns: 'datasetCreation' }) || ''} - onChange={(e) => setInputValue(e.target.value)} - /> -
-
- - -
+ ) diff --git a/web/app/components/datasets/documents/components/__tests__/rename-modal.spec.tsx b/web/app/components/datasets/documents/components/__tests__/rename-modal.spec.tsx index 46db2aa5ed8..f6fe6688041 100644 --- a/web/app/components/datasets/documents/components/__tests__/rename-modal.spec.tsx +++ b/web/app/components/datasets/documents/components/__tests__/rename-modal.spec.tsx @@ -1,4 +1,5 @@ import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import userEvent from '@testing-library/user-event' import { beforeEach, describe, expect, it, vi } from 'vite-plus/test' // Import after mock import { renameDocumentName } from '@/service/datasets' @@ -44,7 +45,7 @@ describe('RenameModal', () => { it('should render name label', () => { render() - expect(screen.getByText(/list\.table\.name/i)).toBeInTheDocument() + expect(screen.getByRole('textbox', { name: /list\.table\.name/i })).toBeInTheDocument() }) it('should render input with initial name', () => { @@ -130,7 +131,7 @@ describe('RenameModal', () => { }) describe('Loading State', () => { - it('should show loading state while saving', async () => { + it('should not submit again while saving', async () => { // Create a promise that we can resolve manually let resolvePromise: (value: { result: 'success' | 'fail' }) => void const pendingPromise = new Promise<{ result: 'success' | 'fail' }>((resolve) => { @@ -139,18 +140,21 @@ describe('RenameModal', () => { mockRenameDocumentName.mockReturnValueOnce(pendingPromise) render() - const saveButton = screen.getByText(/operation\.save/i) - fireEvent.click(saveButton) + const user = userEvent.setup() + const input = screen.getByRole('textbox', { name: /list\.table\.name/i }) + await user.click(input) + await user.keyboard('{Enter}') - // The button should be in loading state await waitFor(() => { - const buttons = screen.getAllByRole('button') - const saveBtn = buttons.find((btn) => btn.textContent?.includes('operation.save')) - expect(saveBtn).toBeInTheDocument() + expect(mockRenameDocumentName).toHaveBeenCalledTimes(1) }) + await user.keyboard('{Enter}') + expect(mockRenameDocumentName).toHaveBeenCalledTimes(1) - // Resolve the promise to clean up resolvePromise!({ result: 'success' }) + await waitFor(() => { + expect(defaultProps.onClose).toHaveBeenCalledTimes(1) + }) }) }) @@ -180,6 +184,7 @@ describe('RenameModal', () => { render() const input = screen.getByRole('textbox') expect(input).toHaveValue('') + expect(input).toHaveAttribute('placeholder', 'common.placeholder.input') }) it('should handle name with special characters', () => { diff --git a/web/app/components/datasets/documents/components/rename-modal.tsx b/web/app/components/datasets/documents/components/rename-modal.tsx index 6a4e9ba6f19..1875e23f385 100644 --- a/web/app/components/datasets/documents/components/rename-modal.tsx +++ b/web/app/components/datasets/documents/components/rename-modal.tsx @@ -2,12 +2,14 @@ import type { FC } from 'react' import { Button } from '@langgenius/dify-ui/button' import { Dialog, DialogContent, DialogTitle } from '@langgenius/dify-ui/dialog' +import { Field, FieldLabel } from '@langgenius/dify-ui/field' +import { Form } from '@langgenius/dify-ui/form' +import { Input } from '@langgenius/dify-ui/input' import { toast } from '@langgenius/dify-ui/toast' import { useBoolean } from 'ahooks' import * as React from 'react' import { useState } from 'react' import { useTranslation } from 'react-i18next' -import Input from '@/app/components/base/input' import { renameDocumentName } from '@/service/datasets' type Props = Readonly<{ @@ -26,6 +28,8 @@ const RenameModal: FC = ({ documentId, datasetId, name, onClose, onSaved useBoolean(false) const handleSave = async () => { + if (saveLoading) return + setSaveLoadingTrue() try { await renameDocumentName({ @@ -54,20 +58,28 @@ const RenameModal: FC = ({ documentId, datasetId, name, onClose, onSaved {t(($) => $['list.table.rename'], { ns: 'datasetDocuments' })} +
void handleSave()}> + + + {t(($) => $['list.table.name'], { ns: 'datasetDocuments' })} + + $['placeholder.input'], { ns: 'common' }) || ''} + onValueChange={setNewName} + /> + -
- {t(($) => $['list.table.name'], { ns: 'datasetDocuments' })} -
- setNewName(e.target.value)} /> - -
- - -
+
+ + +
+
) diff --git a/web/app/components/datasets/list/__tests__/index.spec.tsx b/web/app/components/datasets/list/__tests__/index.spec.tsx index d7966d8680b..c1c7383f5de 100644 --- a/web/app/components/datasets/list/__tests__/index.spec.tsx +++ b/web/app/components/datasets/list/__tests__/index.spec.tsx @@ -2,6 +2,7 @@ import type { ReactElement, ReactNode } from 'react' import { act, fireEvent, screen, waitFor, within } from '@testing-library/react' import userEvent from '@testing-library/user-event' import { createStore, Provider } from 'jotai' +import { queryClientAtom } from 'jotai-tanstack-query' import { hydrateRoot } from 'react-dom/client' import { renderToString } from 'react-dom/server' import { beforeEach, describe, expect, it, vi } from 'vite-plus/test' @@ -399,14 +400,15 @@ describe('List', () => { await user.click(within(guide).getByRole('button', { name: 'dataset.newKnowledge.gotIt' })) firstRender.unmount() - const store = createStore() - seedRegisteredConsoleStateFixture(store) const { wrapper: NuqsWrapper } = createNuqsTestWrapper() - const { wrapper: QueryWrapper } = createConsoleQueryWrapper({ + const { queryClient, wrapper: QueryWrapper } = createConsoleQueryWrapper({ systemFeatures: { knowledge_fs_enabled: mockConsoleState.knowledgeFsEnabled, }, }) + const store = createStore() + store.set(queryClientAtom, queryClient) + seedRegisteredConsoleStateFixture(store) const app = ( diff --git a/web/app/components/goto-anything/actions/commands/__tests__/direct-commands.spec.ts b/web/app/components/goto-anything/actions/commands/__tests__/direct-commands.spec.ts index bb10791912b..1d9fab38609 100644 --- a/web/app/components/goto-anything/actions/commands/__tests__/direct-commands.spec.ts +++ b/web/app/components/goto-anything/actions/commands/__tests__/direct-commands.spec.ts @@ -23,20 +23,15 @@ vi.mock('react-i18next', async () => { } }) -vi.mock('@/context/i18n', () => ({ - defaultDocBaseUrl: 'https://docs.dify.ai', - getDocHomePath: () => '/home', -})) - -vi.mock('@/i18n-config/language', () => ({ - getDocLanguage: (locale: string) => (locale === 'en' ? 'en' : locale), -})) - describe('docsCommand', () => { beforeEach(() => { vi.clearAllMocks() }) + afterEach(() => { + docsCommand.unregister?.() + }) + it('has correct metadata', () => { expect(docsCommand.name).toBe('docs') expect(docsCommand.mode).toBe('direct') @@ -45,6 +40,7 @@ describe('docsCommand', () => { it('execute opens documentation in new tab', () => { const openSpy = vi.spyOn(window, 'open').mockImplementation(() => null) + docsCommand.register?.({ getDocsHomeUrl: () => 'https://docs.dify.ai/en/home' }) docsCommand.execute?.() @@ -56,6 +52,22 @@ describe('docsCommand', () => { openSpy.mockRestore() }) + it('execute uses the documentation URL registered by the provider', () => { + const openSpy = vi.spyOn(window, 'open').mockImplementation(() => null) + docsCommand.register?.({ + getDocsHomeUrl: () => 'https://enterprise-docs.dify.ai/en/', + }) + + docsCommand.execute?.() + + expect(openSpy).toHaveBeenCalledWith( + 'https://enterprise-docs.dify.ai/en/', + '_blank', + 'noopener,noreferrer', + ) + openSpy.mockRestore() + }) + it('search returns a single doc result', async () => { const results = await docsCommand.search('', 'en') @@ -77,18 +89,20 @@ describe('docsCommand', () => { }) it('registers navigation.doc command', () => { - docsCommand.register?.({} as Record) + docsCommand.register?.({ getDocsHomeUrl: () => 'https://docs.dify.ai/en/home' }) expect(registerCommands).toHaveBeenCalledWith({ 'navigation.doc': expect.any(Function) }) }) it('registered handler opens doc URL with correct locale', async () => { - docsCommand.register?.({} as Record) + docsCommand.register?.({ + getDocsHomeUrl: () => 'https://enterprise-docs.dify.ai/en/', + }) const openSpy = vi.spyOn(window, 'open').mockImplementation(() => null) const handlers = vi.mocked(registerCommands).mock.calls[0]![0] await handlers['navigation.doc']!() expect(openSpy).toHaveBeenCalledWith( - 'https://docs.dify.ai/en/home', + 'https://enterprise-docs.dify.ai/en/', '_blank', 'noopener,noreferrer', ) diff --git a/web/app/components/goto-anything/actions/commands/__tests__/slash.spec.tsx b/web/app/components/goto-anything/actions/commands/__tests__/slash.spec.tsx index 57cb861897b..bb6cd6a091e 100644 --- a/web/app/components/goto-anything/actions/commands/__tests__/slash.spec.tsx +++ b/web/app/components/goto-anything/actions/commands/__tests__/slash.spec.tsx @@ -181,4 +181,15 @@ describe('SlashCommandProvider', () => { 'refine', ]) }) + + it('should register the enterprise documentation home URL', () => { + const { unmount } = render(, { + systemFeatures: { deployment_edition: 'ENTERPRISE' }, + }) + const docsRegistration = mockRegister.mock.calls.find((call) => call[0].name === 'docs') + + expect(docsRegistration?.[1].getDocsHomeUrl()).toBe('https://enterprise-docs.dify.ai/en/') + + unmount() + }) }) diff --git a/web/app/components/goto-anything/actions/commands/docs.tsx b/web/app/components/goto-anything/actions/commands/docs.tsx index b234af8a7ca..83d19cdb866 100644 --- a/web/app/components/goto-anything/actions/commands/docs.tsx +++ b/web/app/components/goto-anything/actions/commands/docs.tsx @@ -1,17 +1,16 @@ import type { SlashCommandHandler } from './types' import { getI18n } from 'react-i18next' -import { defaultDocBaseUrl, getDocHomePath } from '@/context/i18n' -import { getDocLanguage } from '@/i18n-config/language' import { registerCommands, unregisterCommands } from './command-bus' -// Documentation command dependency types - no external dependencies needed -type DocDeps = Record +type DocDeps = { + getDocsHomeUrl: () => string +} -const getDocsHomeUrl = () => { - const i18n = getI18n() - const currentLocale = i18n.language - const docLanguage = getDocLanguage(currentLocale) - return `${defaultDocBaseUrl}/${docLanguage}${getDocHomePath()}` +let getDocsHomeUrl: (() => string) | undefined + +const openDocsHome = () => { + const url = getDocsHomeUrl?.() + if (url) window.open(url, '_blank', 'noopener,noreferrer') } /** @@ -24,7 +23,7 @@ export const docsCommand: SlashCommandHandler = { // Direct execution function execute: () => { - window.open(getDocsHomeUrl(), '_blank', 'noopener,noreferrer') + openDocsHome() }, search(args: string, locale: string = 'en') { @@ -47,15 +46,17 @@ export const docsCommand: SlashCommandHandler = { ] }, - register(_deps: DocDeps) { + register(deps: DocDeps) { + getDocsHomeUrl = deps.getDocsHomeUrl registerCommands({ 'navigation.doc': async (_args) => { - window.open(getDocsHomeUrl(), '_blank', 'noopener,noreferrer') + openDocsHome() }, }) }, unregister() { + getDocsHomeUrl = undefined unregisterCommands(['navigation.doc']) }, } diff --git a/web/app/components/goto-anything/actions/commands/slash-provider.tsx b/web/app/components/goto-anything/actions/commands/slash-provider.tsx index c4b50095384..1eeedf8fccc 100644 --- a/web/app/components/goto-anything/actions/commands/slash-provider.tsx +++ b/web/app/components/goto-anything/actions/commands/slash-provider.tsx @@ -2,6 +2,7 @@ import { useTheme } from 'next-themes' import { useEffect } from 'react' import { ENABLE_FEATURE_PREVIEW } from '@/config' +import { useDocLink } from '@/context/i18n' import { setLocaleOnClient } from '@/i18n-config' import { accountCommand } from './account' import { communityCommand } from './community' @@ -15,6 +16,7 @@ import { slashCommandRegistry } from './registry' import { themeCommand } from './theme' type SlashCommandDeps = { + getDocsHomeUrl: () => string setTheme: (theme: string) => void setLocale: typeof setLocaleOnClient } @@ -25,7 +27,7 @@ const registerSlashCommands = (deps: SlashCommandDeps) => { setLocale: deps.setLocale as (locale: string) => Promise, }) slashCommandRegistry.register(forumCommand, {}) - slashCommandRegistry.register(docsCommand, {}) + slashCommandRegistry.register(docsCommand, { getDocsHomeUrl: deps.getDocsHomeUrl }) slashCommandRegistry.register(communityCommand, {}) slashCommandRegistry.register(accountCommand, {}) slashCommandRegistry.register(goCommand, {}) @@ -49,13 +51,15 @@ const unregisterSlashCommands = () => { export const SlashCommandProvider = () => { const theme = useTheme() + const getDocsHomeUrl = useDocLink() useEffect(() => { registerSlashCommands({ + getDocsHomeUrl, setTheme: theme.setTheme, setLocale: setLocaleOnClient, }) return () => unregisterSlashCommands() - }, [theme.setTheme]) + }, [getDocsHomeUrl, theme.setTheme]) return null } diff --git a/web/app/components/header/account-setting/access-rules-page/permission-set-modal/index.tsx b/web/app/components/header/account-setting/access-rules-page/permission-set-modal/index.tsx index 93c3819e7d0..803573596fb 100644 --- a/web/app/components/header/account-setting/access-rules-page/permission-set-modal/index.tsx +++ b/web/app/components/header/account-setting/access-rules-page/permission-set-modal/index.tsx @@ -13,7 +13,7 @@ import { Input } from '@langgenius/dify-ui/input' import { Textarea } from '@langgenius/dify-ui/textarea' import { useState } from 'react' import { useTranslation } from 'react-i18next' -import { useLocale } from '@/context/i18n' +import { getEnterpriseDocUrl, useLocale } from '@/context/i18n' import { getDocLanguage } from '@/i18n-config/language' import PermissionPicker from './permission-picker' @@ -136,7 +136,7 @@ const PermissionSetModalBody = ({
= {}): Model => ({ ...overrides, }) +const makeProviderSummary = (): ModelProviderSummaryResponse => ({ + provider: 'openai', + plugin_id: 'langgenius/openai', + label: { en_US: 'OpenAI', zh_Hans: 'OpenAI' }, + supported_model_types: ['llm'], + configurate_methods: ['predefined-model'], + preferred_provider_type: 'system', + is_configured: true, + custom_configuration: { + status: 'active', + has_custom_models: false, + available_credentials: [], + current_credential_usable: false, + }, + system_configuration: { enabled: true }, +}) + const renderWithQueryClient = (node: ReactNode) => { const queryClient = createConsoleQueryClient() + queryClient.setQueryData(consoleQuery.workspaces.current.modelProviders.summary.get.key(), { + data: [makeProviderSummary()], + plugins: {}, + }) return render({node}) } diff --git a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-auth-dropdown/credits-exhausted-alert.tsx b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-auth-dropdown/credits-exhausted-alert.tsx index eb17feaca04..6ac640133d7 100644 --- a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-auth-dropdown/credits-exhausted-alert.tsx +++ b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-auth-dropdown/credits-exhausted-alert.tsx @@ -71,7 +71,7 @@ export default function CreditsExhaustedAlert({ {t(($) => $['modelProvider.card.usageLabel'], { ns: 'common' })}
- {/* oxlint-disable-next-line hyoban/prefer-tailwind-icons -- This generated icon class is not available to Tailwind. */} + {/* oxlint-disable-next-line dify/prefer-tailwind-icons -- This generated icon class is not available to Tailwind. */} {formatNumber(usedCredits)}/{formatNumber(totalCredits)} diff --git a/web/app/components/header/account-setting/permissions-page/role-modal/index.tsx b/web/app/components/header/account-setting/permissions-page/role-modal/index.tsx index 6950017f9ed..37254f87e4d 100644 --- a/web/app/components/header/account-setting/permissions-page/role-modal/index.tsx +++ b/web/app/components/header/account-setting/permissions-page/role-modal/index.tsx @@ -13,7 +13,7 @@ import { Input } from '@langgenius/dify-ui/input' import { Textarea } from '@langgenius/dify-ui/textarea' import { useCallback, useState } from 'react' import { useTranslation } from 'react-i18next' -import { useLocale } from '@/context/i18n' +import { getEnterpriseDocUrl, useLocale } from '@/context/i18n' import { getDocLanguage } from '@/i18n-config/language' import PermissionField from './permission-field' @@ -113,7 +113,7 @@ const RoleModal = ({ mode, open, role, onClose, onSubmit }: RoleModalProps) => {
= ({ {hasPluginIcon ? ( ) : ( - // oxlint-disable-next-line hyoban/prefer-tailwind-icons -- Reuse the same MagicBox component as the marketplace install button. + // oxlint-disable-next-line dify/prefer-tailwind-icons -- Reuse the same MagicBox component as the marketplace install button. )}
{statusIcon}
diff --git a/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx b/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx index b9cb9c000cf..d1fe99577cd 100644 --- a/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx +++ b/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx @@ -24,6 +24,7 @@ const mockHandleStartWorkflowRun = vi.fn() const mockHandleStopRun = vi.fn() const mockHandleWorkflowStartRunInWorkflow = vi.fn() const mockHandleCheckBeforePublish = vi.fn() +const mockHandleExportDSL = vi.fn() const mockUseAvailableNodesMetaData = vi.hoisted(() => vi.fn()) const mockConsoleState = vi.hoisted(() => ({ workspacePermissionKeys: ['snippets.create_and_modify'] as string[], @@ -141,6 +142,12 @@ vi.mock('../../hooks/use-snippet-start-run', () => ({ }), })) +vi.mock('../hooks/use-snippet-dsl', () => ({ + useSnippetDSL: () => ({ + handleExportDSL: mockHandleExportDSL, + }), +})) + vi.mock('@/app/components/workflow', () => ({ WorkflowWithInnerContext: ({ children, @@ -642,4 +649,12 @@ describe('SnippetMain', () => { }) }) }) + + describe('DSL Export', () => { + it('should pass the snippet DSL export handler to WorkflowWithInnerContext', () => { + renderSnippetMain() + + expect(capturedHooksStore?.handleExportDSL).toBe(mockHandleExportDSL) + }) + }) }) diff --git a/web/app/components/snippets/components/hooks/__tests__/use-snippet-dsl.spec.ts b/web/app/components/snippets/components/hooks/__tests__/use-snippet-dsl.spec.ts new file mode 100644 index 00000000000..0a02982a096 --- /dev/null +++ b/web/app/components/snippets/components/hooks/__tests__/use-snippet-dsl.spec.ts @@ -0,0 +1,62 @@ +import { toast } from '@langgenius/dify-ui/toast' +import { renderHook } from '@testing-library/react' +import { act } from 'react' +import { useExportSnippetMutation } from '@/service/use-snippets' +import { downloadBlob } from '@/utils/download' +import { useSnippetDSL } from '../use-snippet-dsl' + +const mockMutateAsync = vi.fn() + +vi.mock('@/service/use-snippets', () => ({ + useExportSnippetMutation: vi.fn(() => ({ + mutateAsync: mockMutateAsync, + })), +})) + +vi.mock('@/utils/download', () => ({ + downloadBlob: vi.fn(), +})) + +vi.mock('@langgenius/dify-ui/toast', () => ({ + toast: { + error: vi.fn(), + }, +})) + +describe('useSnippetDSL', () => { + beforeEach(() => { + vi.clearAllMocks() + mockMutateAsync.mockResolvedValue('kind: snippet') + }) + + it('exports the requested historical workflow version', async () => { + const { result } = renderHook(() => + useSnippetDSL({ snippetId: 'snippet-1', snippetName: 'My Snippet' }), + ) + + await act(() => result.current.handleExportDSL(false, 'workflow-1')) + + expect(useExportSnippetMutation).toHaveBeenCalled() + expect(mockMutateAsync).toHaveBeenCalledWith({ + snippetId: 'snippet-1', + include: false, + workflowId: 'workflow-1', + }) + expect(downloadBlob).toHaveBeenCalledWith({ + data: expect.any(Blob), + fileName: 'My Snippet.yml', + }) + }) + + it('shows an error when exporting fails', async () => { + mockMutateAsync.mockRejectedValueOnce(new Error('failed')) + const { result } = renderHook(() => + useSnippetDSL({ snippetId: 'snippet-1', snippetName: 'My Snippet' }), + ) + + await act(() => result.current.handleExportDSL(false, 'workflow-1')) + + expect(toast.error).toHaveBeenCalledWith('snippet.exportFailed') + expect(downloadBlob).not.toHaveBeenCalled() + }) +}) diff --git a/web/app/components/snippets/components/hooks/use-snippet-dsl.ts b/web/app/components/snippets/components/hooks/use-snippet-dsl.ts new file mode 100644 index 00000000000..eba5a0879e9 --- /dev/null +++ b/web/app/components/snippets/components/hooks/use-snippet-dsl.ts @@ -0,0 +1,34 @@ +import { toast } from '@langgenius/dify-ui/toast' +import { useCallback } from 'react' +import { useTranslation } from 'react-i18next' +import { useExportSnippetMutation } from '@/service/use-snippets' +import { downloadBlob } from '@/utils/download' + +type UseSnippetDSLOptions = { + snippetId: string + snippetName: string +} + +export const useSnippetDSL = ({ snippetId, snippetName }: UseSnippetDSLOptions) => { + const { t } = useTranslation('snippet') + const exportSnippetMutation = useExportSnippetMutation() + + const handleExportDSL = useCallback( + async (include = false, workflowId?: string) => { + try { + const data = await exportSnippetMutation.mutateAsync({ + snippetId, + include, + workflowId, + }) + const file = new Blob([data], { type: 'application/yaml' }) + downloadBlob({ data: file, fileName: `${snippetName}.yml` }) + } catch { + toast.error(t(($) => $.exportFailed)) + } + }, + [exportSnippetMutation, snippetId, snippetName, t], + ) + + return { handleExportDSL } +} diff --git a/web/app/components/snippets/components/snippet-main.tsx b/web/app/components/snippets/components/snippet-main.tsx index 58b742e2492..0e15a90fffa 100644 --- a/web/app/components/snippets/components/snippet-main.tsx +++ b/web/app/components/snippets/components/snippet-main.tsx @@ -26,6 +26,7 @@ import { useSnippetRun } from '../hooks/use-snippet-run' import { useSnippetStartRun } from '../hooks/use-snippet-start-run' import { useSnippetDetailStore } from '../store' import { canCreateAndModifySnippets } from '../utils/permission' +import { useSnippetDSL } from './hooks/use-snippet-dsl' import { useSnippetInputFieldActions } from './hooks/use-snippet-input-field-actions' import { useSnippetPublish } from './hooks/use-snippet-publish' import SnippetChildren from './snippet-children' @@ -205,6 +206,7 @@ const SnippetMain = ({ canEdit: canEditSnippet, snippetId, }) + const { handleExportDSL } = useSnippetDSL({ snippetId, snippetName: snippet.name }) const { handleStartWorkflowRun, handleWorkflowStartRunInWorkflow } = useSnippetStartRun({ handleRun, }) @@ -322,6 +324,7 @@ const SnippetMain = ({ handleStopRun, handleStartWorkflowRun, handleWorkflowStartRunInWorkflow, + handleExportDSL, getWorkflowRunAndTraceUrl, availableNodesMetaData, fetchInspectVars, @@ -368,6 +371,7 @@ const SnippetMain = ({ handleStartWorkflowRun, handleStopRun, handleWorkflowStartRunInWorkflow, + handleExportDSL, getWorkflowRunAndTraceUrl, hasNodeInspectVars, hasSetInspectVar, diff --git a/web/app/components/workflow-app/components/__tests__/workflow-main.spec.tsx b/web/app/components/workflow-app/components/__tests__/workflow-main.spec.tsx index 600a6a404d0..c1af965355c 100644 --- a/web/app/components/workflow-app/components/__tests__/workflow-main.spec.tsx +++ b/web/app/components/workflow-app/components/__tests__/workflow-main.spec.tsx @@ -5,6 +5,7 @@ import { useStore as useAppStore } from '@/app/components/app/store' import { ChatVarType } from '@/app/components/workflow/panel/chat-variable-panel/type' import { BlockEnum } from '@/app/components/workflow/types' import { renderWithAccountProfile as render } from '@/test/console/account-profile' +import { AppACLPermission } from '@/utils/permission' import WorkflowMain from '../workflow-main' const mockSetFeatures = vi.fn() @@ -24,6 +25,7 @@ const mockReplaceGraphFromReactFlow = vi.hoisted(() => vi.fn()) const mockCanPersistLocalGraph = vi.hoisted(() => vi.fn()) const mockIsGraphReloadCurrent = vi.hoisted(() => vi.fn()) const mockRetryGraphReload = vi.hoisted(() => vi.fn()) +const mockUseCollaboration = vi.hoisted(() => vi.fn()) const hookFns = { doSyncWorkflowDraft: vi.fn(), @@ -135,14 +137,10 @@ vi.mock('reactflow', () => ({ })) vi.mock('@/app/components/workflow/collaboration/hooks/use-collaboration', () => ({ - useCollaboration: () => ({ - startCursorTracking: collaborationRuntime.startCursorTracking, - stopCursorTracking: collaborationRuntime.stopCursorTracking, - onlineUsers: collaborationRuntime.onlineUsers, - cursors: collaborationRuntime.cursors, - isConnected: collaborationRuntime.isConnected, - isEnabled: collaborationRuntime.isEnabled, - }), + useCollaboration: (...args: unknown[]) => { + mockUseCollaboration(...args) + return collaborationRuntime + }, })) vi.mock('@/app/components/workflow/hooks/use-workflow-update', () => ({ @@ -590,6 +588,16 @@ describe('WorkflowMain', () => { expect(screen.queryByRole('status')).not.toBeInTheDocument() }) + it('disables collaboration for view-only apps', () => { + useAppStore.setState({ + appDetail: { permission_keys: [AppACLPermission.ViewLayout] } as never, + }) + + render() + + expect(mockUseCollaboration).toHaveBeenCalledWith('app-1', false, expect.any(Object)) + }) + it('subscribes collaboration listeners and handles sync/workflow update callbacks', async () => { collaborationRuntime.isEnabled = true mockFetchWorkflowDraft.mockResolvedValue({ diff --git a/web/app/components/workflow-app/components/workflow-main.tsx b/web/app/components/workflow-app/components/workflow-main.tsx index 14bb178ff12..dfa7666aef6 100644 --- a/web/app/components/workflow-app/components/workflow-main.tsx +++ b/web/app/components/workflow-app/components/workflow-main.tsx @@ -74,22 +74,6 @@ const WorkflowMain = ({ nodes, edges, viewport }: WorkflowMainProps) => { }), [reactFlow], ) - const { - startCursorTracking, - stopCursorTracking, - onlineUsers, - cursors, - isConnected, - isEnabled: isCollaborationEnabled, - } = useCollaboration(appId || '', reactFlowStore) - const myUserId = useMemo( - () => (isCollaborationEnabled && isConnected ? 'current-user' : null), - [isCollaborationEnabled, isConnected], - ) - - const filteredCursors = Object.fromEntries( - Object.entries(cursors).filter(([userId]) => userId !== myUserId), - ) const { data: currentUserId } = useSuspenseQuery({ ...userProfileQueryOptions(), select: (data) => data.profile.id, @@ -104,6 +88,22 @@ const WorkflowMain = ({ nodes, edges, viewport }: WorkflowMainProps) => { }), [appDetail?.maintainer, appDetail?.permission_keys, currentUserId, workspacePermissionKeys], ) + const { + startCursorTracking, + stopCursorTracking, + onlineUsers, + cursors, + isConnected, + isEnabled: isCollaborationEnabled, + } = useCollaboration(appId || '', appACLCapabilities.canEdit, reactFlowStore) + const myUserId = useMemo( + () => (isCollaborationEnabled && isConnected ? 'current-user' : null), + [isCollaborationEnabled, isConnected], + ) + + const filteredCursors = Object.fromEntries( + Object.entries(cursors).filter(([userId]) => userId !== myUserId), + ) useEffect(() => { if (!isCollaborationEnabled) return diff --git a/web/app/components/workflow/collaboration/hooks/__tests__/use-collaboration.spec.ts b/web/app/components/workflow/collaboration/hooks/__tests__/use-collaboration.spec.ts index 828d8275e8a..61b079e74ee 100644 --- a/web/app/components/workflow/collaboration/hooks/__tests__/use-collaboration.spec.ts +++ b/web/app/components/workflow/collaboration/hooks/__tests__/use-collaboration.spec.ts @@ -3,7 +3,7 @@ import { waitFor } from '@testing-library/react' import { renderHookWithConsoleQuery } from '@/test/console/query-data' import { useCollaboration } from '../use-collaboration' -type HookReactFlowStore = NonNullable[1]> +type HookReactFlowStore = NonNullable[2]> type HookReactFlowInstance = Parameters< ReturnType['startCursorTracking'] >[1] @@ -105,7 +105,7 @@ describe('useCollaboration', () => { getState: vi.fn(), } const { result, unmount } = renderHookWithConsoleQuery( - () => useCollaboration('app-1', reactFlowStore), + () => useCollaboration('app-1', true, reactFlowStore), { systemFeatures: { enable_collaboration_mode: isCollaborationEnabled }, }, @@ -155,18 +155,24 @@ describe('useCollaboration', () => { expect(mockDisconnect).toHaveBeenCalledWith('conn-1') }) - it('does not connect or start cursor tracking when collaboration is disabled', async () => { - isCollaborationEnabled = false - const { result } = renderHookWithConsoleQuery(() => useCollaboration('app-1'), { - systemFeatures: { enable_collaboration_mode: isCollaborationEnabled }, - }) + it.each([ + [false, true], + [true, false], + ])( + 'does not connect or track cursors when a collaboration gate is disabled', + async (featureEnabled, canEdit) => { + isCollaborationEnabled = featureEnabled + const { result } = renderHookWithConsoleQuery(() => useCollaboration('app-1', canEdit), { + systemFeatures: { enable_collaboration_mode: featureEnabled }, + }) - await waitFor(() => { - expect(mockConnect).not.toHaveBeenCalled() - expect(result.current.isEnabled).toBe(false) - }) + await waitFor(() => { + expect(mockConnect).not.toHaveBeenCalled() + expect(result.current.isEnabled).toBe(false) + }) - result.current.startCursorTracking({ current: document.createElement('div') }) - expect(mockStartTracking).not.toHaveBeenCalled() - }) + result.current.startCursorTracking({ current: document.createElement('div') }) + expect(mockStartTracking).not.toHaveBeenCalled() + }, + ) }) diff --git a/web/app/components/workflow/collaboration/hooks/use-collaboration.ts b/web/app/components/workflow/collaboration/hooks/use-collaboration.ts index 12f6436263b..4d4002999a2 100644 --- a/web/app/components/workflow/collaboration/hooks/use-collaboration.ts +++ b/web/app/components/workflow/collaboration/hooks/use-collaboration.ts @@ -29,7 +29,7 @@ const initialState: CollaborationViewState = { isLeader: false, } -export function useCollaboration(appId: string, reactFlowStore?: ReactFlowStore) { +export function useCollaboration(appId: string, canEdit: boolean, reactFlowStore?: ReactFlowStore) { const [state, setState] = useState(initialState) const cursorServiceRef = useRef(null) @@ -40,7 +40,7 @@ export function useCollaboration(appId: string, reactFlowStore?: ReactFlowStore) }) useEffect(() => { - if (!appId || !isCollaborationEnabled) { + if (!appId || !isCollaborationEnabled || !canEdit) { Promise.resolve().then(() => { setState(initialState) }) @@ -110,7 +110,7 @@ export function useCollaboration(appId: string, reactFlowStore?: ReactFlowStore) cursorServiceRef.current?.stopTracking() if (connectionId) collaborationManager.disconnect(connectionId) } - }, [appId, reactFlowStore, isCollaborationEnabled]) + }, [appId, canEdit, reactFlowStore, isCollaborationEnabled]) const prevIsConnected = useRef(false) useEffect(() => { @@ -126,7 +126,7 @@ export function useCollaboration(appId: string, reactFlowStore?: ReactFlowStore) containerRef: React.RefObject, reactFlowInstance?: ReactFlowInstance, ) => { - if (!isCollaborationEnabled || !cursorServiceRef.current) return + if (!isCollaborationEnabled || !canEdit || !cursorServiceRef.current) return if (cursorServiceRef.current) { cursorServiceRef.current.startTracking( @@ -150,7 +150,7 @@ export function useCollaboration(appId: string, reactFlowStore?: ReactFlowStore) nodePanelPresence: state.nodePanelPresence || {}, isLeader: state.isLeader || false, leaderId: collaborationManager.getLeaderId(), - isEnabled: isCollaborationEnabled, + isEnabled: isCollaborationEnabled && canEdit, startCursorTracking, stopCursorTracking, } diff --git a/web/app/components/workflow/header/online-users.tsx b/web/app/components/workflow/header/online-users.tsx index 427b9423e63..580ddaaefd9 100644 --- a/web/app/components/workflow/header/online-users.tsx +++ b/web/app/components/workflow/header/online-users.tsx @@ -14,6 +14,7 @@ import { userProfileQueryOptions } from '@/features/account-profile/client' import { getAvatar } from '@/service/common' import { useCollaboration } from '../collaboration/hooks/use-collaboration' import { getUserColor } from '../collaboration/utils/user-color' +import { useHooksStore } from '../hooks-store' import { useStore } from '../store' const useAvatarUrls = (users: OnlineUser[]) => { @@ -49,11 +50,12 @@ const useAvatarUrls = (users: OnlineUser[]) => { const OnlineUsers = () => { const { t } = useTranslation() const appId = useStore((s) => s.appId) + const canEdit = useHooksStore((s) => s.accessControl.canEdit) const { onlineUsers, cursors, isEnabled: isCollaborationEnabled, - } = useCollaboration(appId as string) + } = useCollaboration(appId as string, canEdit) const { data: currentUserId } = useSuspenseQuery({ ...userProfileQueryOptions(), select: (data) => data.profile.id, @@ -119,29 +121,43 @@ const OnlineUsers = () => { const userColor = isCurrentUser ? undefined : getUserColor(user.user_id) const avatarUrl = getAvatarUrl(user) const displayName = user.username || fallbackUsername + const avatar = ( + + {avatarUrl && } + + {displayName?.[0]?.toLocaleUpperCase()} + + + ) + const triggerClassName = cn( + 'relative flex size-6 items-center justify-center', + index > 0 && '-ml-1.5', + !isCurrentUser && 'cursor-pointer transition-transform hover:scale-110', + ) + const triggerStyle = { zIndex: visibleUsers.length - index } return ( - - -
0 && '-ml-1.5', - !isCurrentUser && 'cursor-pointer transition-transform hover:scale-110', - )} - style={{ zIndex: visibleUsers.length - index }} - onClick={() => !isCurrentUser && jumpToUserCursor(user.user_id)} - > - - {avatarUrl && } - + + {avatar} +
+ ) : ( +
- + {avatar} + + ) + } + /> { const avatarUrl = getAvatarUrl(user) const displayName = user.username || fallbackUsername return ( -
{ - if (!isCurrentUser) { - jumpToUserCursor(user.user_id) - setDropdownOpen(false) - } + jumpToUserCursor(user.user_id) + setDropdownOpen(false) }} >
@@ -219,7 +235,7 @@ const OnlineUsers = () => { 'system-xs-medium text-text-secondary', 'text-text-tertiary', )} -
+ ) })} diff --git a/web/app/components/workflow/nodes/_base/components/workflow-panel/index.tsx b/web/app/components/workflow/nodes/_base/components/workflow-panel/index.tsx index 8956c99e41c..2c035f07b2b 100644 --- a/web/app/components/workflow/nodes/_base/components/workflow-panel/index.tsx +++ b/web/app/components/workflow/nodes/_base/components/workflow-panel/index.tsx @@ -100,7 +100,8 @@ const BasePanel: FC = ({ id, data, children }) => { ...userProfileQueryOptions(), select: (data) => data.profile, }) - const { isConnected, nodePanelPresence } = useCollaboration(appId as string) + const canEdit = useHooksStore((s) => s.accessControl.canEdit) + const { isConnected, nodePanelPresence } = useCollaboration(appId as string, canEdit) const { showMessageLogModal } = useAppStore( useShallow((state) => ({ showMessageLogModal: state.showMessageLogModal, diff --git a/web/app/components/workflow/nodes/_base/node.tsx b/web/app/components/workflow/nodes/_base/node.tsx index 5990ddad373..61b255cd095 100644 --- a/web/app/components/workflow/nodes/_base/node.tsx +++ b/web/app/components/workflow/nodes/_base/node.tsx @@ -9,6 +9,7 @@ import { UserAvatarList } from '@/app/components/base/user-avatar-list' import BlockIcon from '@/app/components/workflow/block-icon' import { ToolType } from '@/app/components/workflow/block-selector/types' import { useCollaboration } from '@/app/components/workflow/collaboration/hooks/use-collaboration' +import { useHooksStore } from '@/app/components/workflow/hooks-store' import { useNodeIterationInteractions } from '@/app/components/workflow/nodes/iteration/use-interactions' import { useNodeLoopInteractions } from '@/app/components/workflow/nodes/loop/use-interactions' import CopyID from '@/app/components/workflow/nodes/tool/components/copy-id' @@ -62,7 +63,8 @@ const BaseNode: FC = ({ id, data, children }) => { select: (data) => data.profile, }) const appId = useStore((s) => s.appId) - const { nodePanelPresence } = useCollaboration(appId as string) + const canEdit = useHooksStore((s) => s.accessControl.canEdit) + const { nodePanelPresence } = useCollaboration(appId as string, canEdit) const controlMode = useStore((s) => s.controlMode) const isContextMenuTarget = useStore( (s) => s.contextMenuTarget?.type === 'node' && s.contextMenuTarget.nodeId === id, diff --git a/web/app/components/workflow/nodes/trigger-webhook/components/generic-table.tsx b/web/app/components/workflow/nodes/trigger-webhook/components/generic-table.tsx index a33add66a9b..cbe2783d5e8 100644 --- a/web/app/components/workflow/nodes/trigger-webhook/components/generic-table.tsx +++ b/web/app/components/workflow/nodes/trigger-webhook/components/generic-table.tsx @@ -325,7 +325,7 @@ const GenericTable: FC = ({ className="p-1" aria-label="Delete row" > - {/* oxlint-disable-next-line hyoban/prefer-tailwind-icons */} + {/* oxlint-disable-next-line dify/prefer-tailwind-icons */}
diff --git a/web/app/components/workflow/panel/env-panel/index.tsx b/web/app/components/workflow/panel/env-panel/index.tsx index 94ed5257791..89ed16c100b 100644 --- a/web/app/components/workflow/panel/env-panel/index.tsx +++ b/web/app/components/workflow/panel/env-panel/index.tsx @@ -802,7 +802,7 @@ const EnvPanel = () => { className="flex size-6 cursor-pointer items-center justify-center" onClick={() => setShowEnvPanel(false)} > - {/* oxlint-disable-next-line hyoban/prefer-tailwind-icons */} + {/* oxlint-disable-next-line dify/prefer-tailwind-icons */}
diff --git a/web/app/components/workflow/run/iteration-log/iteration-log-trigger.tsx b/web/app/components/workflow/run/iteration-log/iteration-log-trigger.tsx index e459a284958..1020a90e9d6 100644 --- a/web/app/components/workflow/run/iteration-log/iteration-log-trigger.tsx +++ b/web/app/components/workflow/run/iteration-log/iteration-log-trigger.tsx @@ -135,7 +135,7 @@ const IterationLogTrigger = ({ className="flex w-full cursor-pointer items-center self-stretch rounded-lg bg-components-button-tertiary-bg-hover px-3 py-2 inset-ring-0 hover:bg-components-button-tertiary-bg-hover" onClick={handleOnShowIterationDetail} > - {/* oxlint-disable-next-line hyoban/prefer-tailwind-icons */} + {/* oxlint-disable-next-line dify/prefer-tailwind-icons */}
{t(($) => $['nodes.iteration.iteration'], { ns: 'workflow', count: displayIterationCount })} @@ -146,7 +146,7 @@ const IterationLogTrigger = ({ )}
- {/* oxlint-disable-next-line hyoban/prefer-tailwind-icons */} + {/* oxlint-disable-next-line dify/prefer-tailwind-icons */} ) diff --git a/web/app/device/__tests__/page-terminal.spec.tsx b/web/app/device/__tests__/page-terminal.spec.tsx index 3d74e85bd8f..8262d9dc267 100644 --- a/web/app/device/__tests__/page-terminal.spec.tsx +++ b/web/app/device/__tests__/page-terminal.spec.tsx @@ -1,5 +1,5 @@ import { useQuery } from '@tanstack/react-query' -import { fireEvent, screen } from '@testing-library/react' +import { fireEvent, screen, waitFor } from '@testing-library/react' import { beforeEach, describe, expect, it, vi } from 'vite-plus/test' import { renderWithConsoleQuery as render } from '@/test/console/query-data' import DevicePage from '../page' @@ -80,7 +80,9 @@ describe('error_expired terminal state', () => { it('shows "errorExpired.title" heading', async () => { await reachTerminal(new Error('expired')) await screen.findByText('deviceFlow.errorExpired.title') - expect(document.title).toBe('deviceFlow.errorExpired.title - Dify') + await waitFor(() => { + expect(document.title).toBe('deviceFlow.errorExpired.title - Dify') + }) }) it('ghost button resets to code_entry', async () => { diff --git a/web/app/install/installForm.spec.tsx b/web/app/install/installForm.spec.tsx index 7d2db250216..32904d60d05 100644 --- a/web/app/install/installForm.spec.tsx +++ b/web/app/install/installForm.spec.tsx @@ -1,6 +1,7 @@ import type { ReactElement } from 'react' import type { InitValidateStatusResponse, SetupStatusResponse } from '@/models/common' import { fireEvent, screen, waitFor } from '@testing-library/react' +import userEvent from '@testing-library/user-event' import { fetchInitValidateStatus, fetchSetupStatus, login, setup } from '@/service/common' import { expectLoadingButton } from '@/test/button' import { renderWithConsoleQuery } from '@/test/console/query-data' @@ -44,22 +45,85 @@ describe('InstallForm', () => { it('should render form after loading', async () => { render() - expect(await screen.findByLabelText('login.email')).toBeInTheDocument() + const emailInput = await screen.findByLabelText('login.email') + const nameInput = screen.getByLabelText('login.name') + const passwordInput = screen.getByLabelText('login.password') + + expect(emailInput).toHaveAttribute('type', 'email') + expect(emailInput).toHaveAttribute('autocomplete', 'email') + expect(nameInput).toHaveAttribute('autocomplete', 'name') + expect(nameInput).toHaveAttribute('maxlength', '30') + expect(passwordInput).toHaveAttribute('autocomplete', 'new-password') expect(screen.getByRole('heading', { level: 1 })).toHaveTextContent('login.setAdminAccount') expect(screen.getByRole('button', { name: /login\.installBtn/ })).toBeInTheDocument() }) - it('should show validation error when required fields are empty', async () => { + it('should reveal and hide the password with an accessible action', async () => { + const user = userEvent.setup() render() - await screen.findByLabelText('login.email') + const passwordInput = await screen.findByLabelText('login.password') - fireEvent.click(screen.getByRole('button', { name: /login\.installBtn/ })) + await user.click(screen.getByRole('button', { name: 'login.showPassword' })) + expect(passwordInput).toHaveAttribute('type', 'text') + + await user.click(screen.getByRole('button', { name: 'login.hidePassword' })) + expect(passwordInput).toHaveAttribute('type', 'password') + }) + + it('should identify required fields only after submission', async () => { + const user = userEvent.setup() + render() + + const emailInput = await screen.findByLabelText('login.email') + const nameInput = screen.getByLabelText('login.name') + const passwordInput = screen.getByLabelText('login.password') + + expect(screen.queryByText('login.error.emailInValid')).not.toBeInTheDocument() + expect(screen.queryByText('login.error.nameEmpty')).not.toBeInTheDocument() + expect(screen.getAllByText('login.error.passwordInvalid')).toHaveLength(1) + expect(passwordInput).toHaveAccessibleDescription('login.error.passwordInvalid') + + await user.click(screen.getByRole('button', { name: /login\.installBtn/ })) await waitFor(() => { expect(screen.getByText('login.error.emailInValid')).toBeInTheDocument() expect(screen.getByText('login.error.nameEmpty')).toBeInTheDocument() + expect(screen.getAllByText('login.error.passwordInvalid')).toHaveLength(1) }) + expect(emailInput).toHaveAttribute('aria-invalid', 'true') + expect(nameInput).toHaveAttribute('aria-invalid', 'true') + expect(passwordInput).toHaveAttribute('aria-invalid', 'true') + expect(passwordInput).toHaveAccessibleDescription('login.error.passwordInvalid') + expect(mockSetup).not.toHaveBeenCalled() + }) + + it('should identify an invalid email and focus the field', async () => { + const user = userEvent.setup() + render() + + const emailInput = await screen.findByLabelText('login.email') + await user.type(emailInput, 'invalid-email') + await user.click(screen.getByRole('button', { name: /login\.installBtn/ })) + + expect(await screen.findByText('login.error.emailInValid')).toBeInTheDocument() + expect(emailInput).toHaveAttribute('aria-invalid', 'true') + expect(emailInput).toHaveFocus() + expect(mockSetup).not.toHaveBeenCalled() + }) + + it('should enforce the password requirements before submission', async () => { + const user = userEvent.setup() + render() + + await user.type(await screen.findByLabelText('login.email'), 'admin@example.com') + await user.type(screen.getByLabelText('login.name'), 'Admin') + const passwordInput = screen.getByLabelText('login.password') + await user.type(passwordInput, 'abcdefgh') + await user.click(screen.getByRole('button', { name: /login\.installBtn/ })) + + expect(passwordInput).toHaveAttribute('aria-invalid', 'true') + expect(passwordInput).toHaveFocus() expect(mockSetup).not.toHaveBeenCalled() }) diff --git a/web/app/install/installForm.tsx b/web/app/install/installForm.tsx index d11a64206b8..738f8a68ac5 100644 --- a/web/app/install/installForm.tsx +++ b/web/app/install/installForm.tsx @@ -1,16 +1,23 @@ 'use client' import type { InitValidateStatusResponse, SetupStatusResponse } from '@/models/common' +import { zPostSetupBody } from '@dify/contracts/api/console/setup/zod.gen' import { Button } from '@langgenius/dify-ui/button' -import { cn } from '@langgenius/dify-ui/cn' -import { useStore } from '@tanstack/react-form' +import { + Field, + FieldDescription, + FieldError, + FieldLabel, + FieldValidity, +} from '@langgenius/dify-ui/field' +import { Form } from '@langgenius/dify-ui/form' +import { IconButton } from '@langgenius/dify-ui/icon-button' +import { Input } from '@langgenius/dify-ui/input' +import { InputGroup, InputGroupAddon, InputGroupInput } from '@langgenius/dify-ui/input-group' import { useQueryClient } from '@tanstack/react-query' import * as React from 'react' import { useEffect } from 'react' import { useTranslation } from 'react-i18next' import * as z from 'zod' -import { formContext, useAppForm } from '@/app/components/base/form' -import { zodSubmitValidator } from '@/app/components/base/form/utils/zod-submit-validator' -import Input from '@/app/components/base/input' import { validPassword } from '@/config' import { LICENSE_LINK } from '@/constants/link' import useDocumentTitle from '@/hooks/use-document-title' @@ -21,40 +28,29 @@ import { fetchInitValidateStatus, fetchSetupStatus, login, setup } from '@/servi import { encryptPassword as encodePassword } from '@/utils/encryption' import Loading from '../components/base/loading' -const accountFormSchema = z.object({ - email: z.email('error.emailInValid').min(1, { - error: 'error.emailInValid', - }), - name: z.string().min(1, { - error: 'error.nameEmpty', - }), - password: z - .string() - .min(8, { - error: 'error.passwordLengthInValid', - }) - .regex(validPassword, 'error.passwordInvalid'), +const accountFormSchema = zPostSetupBody.pick({ email: true, name: true, password: true }).extend({ + email: zPostSetupBody.shape.email.pipe(z.email()), + name: zPostSetupBody.shape.name.min(1), + password: zPostSetupBody.shape.password.min(8).regex(validPassword), }) +type AccountFormValues = z.infer + const InstallForm = () => { const { t, i18n } = useTranslation() const pageTitle = t(($) => $.setAdminAccount, { ns: 'login' }) useDocumentTitle(pageTitle) - const router = useRouter() + const { push, replace } = useRouter() const queryClient = useQueryClient() const [showPassword, setShowPassword] = React.useState(false) const [loading, setLoading] = React.useState(true) + const [isSubmitting, setIsSubmitting] = React.useState(false) - const form = useAppForm({ - defaultValues: { - name: '', - password: '', - email: '', - }, - validators: { - onSubmit: zodSubmitValidator(accountFormSchema), - }, - onSubmit: async ({ value }) => { + const handleSubmit = async (value: AccountFormValues) => { + if (isSubmitting) return + + setIsSubmitting(true) + try { // First, setup the admin account await setup({ body: { @@ -75,31 +71,28 @@ const InstallForm = () => { // Store tokens and redirect if login successful if (loginRes.result === 'success') { await queryClient.resetQueries({ queryKey: consoleQuery.account.profile.get.key() }) - router.replace('/') + replace('/') } else { // Fallback to signin page if auto-login fails - router.replace('/signin') + replace('/signin') } - }, - }) - - const isSubmitting = useStore(form.store, (state) => state.isSubmitting) - const emailErrors = useStore(form.store, (state) => state.fieldMeta.email?.errors) - const nameErrors = useStore(form.store, (state) => state.fieldMeta.name?.errors) - const passwordErrors = useStore(form.store, (state) => state.fieldMeta.password?.errors) + } finally { + setIsSubmitting(false) + } + } useEffect(() => { fetchSetupStatus().then((res: SetupStatusResponse) => { if (res.step === 'finished') { - router.push('/signin') + push('/signin') } else { fetchInitValidateStatus().then((res: InitValidateStatusResponse) => { - if (res.status === 'not_started') router.push('/init') + if (res.status === 'not_started') push('/init') }) } setLoading(false) }) - }, []) + }, [push]) return loading ? ( @@ -113,123 +106,97 @@ const InstallForm = () => {
- -
{ - e.preventDefault() - e.stopPropagation() - if (isSubmitting) return - form.handleSubmit() - }} + onFormSubmit={(value) => void handleSubmit(value)}> + + accountFormSchema.shape.email.safeParse(value).success + ? null + : t(($) => $['error.emailInValid'], { ns: 'login' }) + } + className="mb-5" > -
- -
- - {(field) => ( - field.handleChange(e.target.value)} - onBlur={field.handleBlur} - placeholder={t(($) => $.emailPlaceholder, { ns: 'login' }) || ''} - /> - )} - - {emailErrors && emailErrors.length > 0 && ( - - {t(($) => $[`${emailErrors[0]}` as 'error.emailInValid'], { ns: 'login' })} - - )} -
-
+ {t(($) => $.email, { ns: 'login' })} + $.emailPlaceholder, { ns: 'login' }) || ''} + /> + {t(($) => $['error.emailInValid'], { ns: 'login' })} + -
- -
- - {(field) => ( - field.handleChange(e.target.value)} - onBlur={field.handleBlur} - placeholder={t(($) => $.namePlaceholder, { ns: 'login' }) || ''} - /> - )} - -
- {nameErrors && nameErrors.length > 0 && ( - - {t(($) => $[`${nameErrors[0]}` as 'error.nameEmpty'], { ns: 'login' })} - - )} -
+ + accountFormSchema.shape.name.safeParse(value).success + ? null + : t(($) => $['error.nameEmpty'], { ns: 'login' }) + } + className="mb-5" + > + {t(($) => $.name, { ns: 'login' })} + $.namePlaceholder, { ns: 'login' }) || ''} + /> + {t(($) => $['error.nameEmpty'], { ns: 'login' })} + -
- -
- - {(field) => ( - field.handleChange(e.target.value)} - onBlur={field.handleBlur} - placeholder={t(($) => $.passwordPlaceholder, { ns: 'login' }) || ''} - /> - )} - + + accountFormSchema.shape.password.safeParse(value).success + ? null + : t(($) => $['error.passwordInvalid'], { ns: 'login' }) + } + className="mb-5" + > + {t(($) => $.password, { ns: 'login' })} + + $.passwordPlaceholder, { ns: 'login' }) || ''} + /> + + $[showPassword ? 'hidePassword' : 'showPassword'], { + ns: 'login', + })} + onClick={() => setShowPassword((visible) => !visible)} + > + + + + + {({ validity }) => + validity.valid === false ? null : ( + + {t(($) => $['error.passwordInvalid'], { ns: 'login' })} + + ) + } + + {t(($) => $['error.passwordInvalid'], { ns: 'login' })} + -
- -
-
- -
0, - })} - > - {t(($) => $['error.passwordInvalid'], { ns: 'login' })} -
-
- -
- -
- -
+
+ +
+
{t(($) => $['license.tip'], { ns: 'login' })}   diff --git a/web/app/signup/components/input-mail.spec.tsx b/web/app/signup/components/input-mail.spec.tsx index d93c6f26a24..cc517e386c6 100644 --- a/web/app/signup/components/input-mail.spec.tsx +++ b/web/app/signup/components/input-mail.spec.tsx @@ -6,7 +6,7 @@ import { useLocale } from '@/context/i18n' import { useSearchParams } from '@/next/navigation' import { useSendMail } from '@/service/use-common' import { renderWithConsoleQuery } from '@/test/console/query-data' -import Form from './input-mail' +import SignupEmailForm from './input-mail' const mockSubmitMail = vi.fn() const mockOnSuccess = vi.fn() @@ -61,7 +61,7 @@ const renderForm = ({ mutateAsync: mockSubmitMail, isPending, } as unknown as UseSendMailResult) - return renderWithConsoleQuery(
, { + return renderWithConsoleQuery(, { systemFeatures: { branding: { enabled: brandingEnabled } }, }) } diff --git a/web/app/signup/components/input-mail.tsx b/web/app/signup/components/input-mail.tsx index 36af8562859..8312e837747 100644 --- a/web/app/signup/components/input-mail.tsx +++ b/web/app/signup/components/input-mail.tsx @@ -1,11 +1,13 @@ 'use client' import type { MailSendResponse } from '@/service/use-common' import { Button } from '@langgenius/dify-ui/button' +import { Field, FieldLabel } from '@langgenius/dify-ui/field' +import { Form } from '@langgenius/dify-ui/form' +import { Input } from '@langgenius/dify-ui/input' import { toast } from '@langgenius/dify-ui/toast' import { useSuspenseQuery } from '@tanstack/react-query' import { useCallback, useState } from 'react' import { useTranslation } from 'react-i18next' -import Input from '@/app/components/base/input' import Split from '@/app/signin/split' import { emailRegex } from '@/config' import { useLocale } from '@/context/i18n' @@ -17,7 +19,7 @@ import { useSendMail } from '@/service/use-common' type Props = { onSuccess: (email: string, payload: string) => void } -export default function Form({ onSuccess }: Props) { +export default function SignupEmailForm({ onSuccess }: Props) { const { t } = useTranslation() const [email, setEmail] = useState('') const locale = useLocale() @@ -45,29 +47,25 @@ export default function Form({ onSuccess }: Props) { }, [email, locale, submitMail, t, isPending, onSuccess]) return ( - { e.preventDefault() handleSubmit() }} > -
- -
- setEmail(e.target.value)} - id="email" - name="email" - type="email" - autoComplete="email" - spellCheck={false} - placeholder={t(($) => $.emailPlaceholder, { ns: 'login' }) || ''} - /> -
-
+ + $.emailPlaceholder, { ns: 'login' }) || ''} + /> +
)} - + ) } diff --git a/web/app/signup/set-password/__tests__/page.spec.tsx b/web/app/signup/set-password/__tests__/page.spec.tsx index 50021c10254..02bff2877f0 100644 --- a/web/app/signup/set-password/__tests__/page.spec.tsx +++ b/web/app/signup/set-password/__tests__/page.spec.tsx @@ -2,6 +2,7 @@ import type { ReactElement } from 'react' import type { MockedFunction } from 'vite-plus/test' import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import userEvent from '@testing-library/user-event' import Cookies from 'js-cookie' import { beforeEach, describe, expect, it, vi } from 'vite-plus/test' import { useLocale } from '@/context/i18n' @@ -100,15 +101,17 @@ describe('Signup Set Password Page', () => { describe('Registration payload', () => { it('should submit locale and browser timezone when setting password', async () => { + const user = userEvent.setup() renderWithQueryClient() - fireEvent.change(screen.getByLabelText('common.account.newPassword'), { - target: { value: 'ValidPass123!' }, - }) - fireEvent.change(screen.getByLabelText('common.account.confirmPassword'), { - target: { value: 'ValidPass123!' }, - }) - fireEvent.click(screen.getByRole('button', { name: 'login.changePasswordBtn' })) + const passwordInput = screen.getByLabelText('common.account.newPassword') + const confirmPasswordInput = screen.getByLabelText('common.account.confirmPassword') + + expect(passwordInput).toHaveAttribute('autocomplete', 'new-password') + expect(confirmPasswordInput).toHaveAttribute('autocomplete', 'new-password') + + await user.type(passwordInput, 'ValidPass123!') + await user.type(confirmPasswordInput, 'ValidPass123!{Enter}') await waitFor(() => { expect(mockRegister).toHaveBeenCalledWith({ diff --git a/web/app/signup/set-password/page.tsx b/web/app/signup/set-password/page.tsx index 30ec554db3e..c9d9d07992c 100644 --- a/web/app/signup/set-password/page.tsx +++ b/web/app/signup/set-password/page.tsx @@ -2,13 +2,15 @@ import type { MailRegisterResponse } from '@/service/use-common' import { Button } from '@langgenius/dify-ui/button' import { cn } from '@langgenius/dify-ui/cn' +import { Field, FieldDescription, FieldLabel } from '@langgenius/dify-ui/field' +import { Form } from '@langgenius/dify-ui/form' +import { Input } from '@langgenius/dify-ui/input' import { toast } from '@langgenius/dify-ui/toast' import { useQueryClient } from '@tanstack/react-query' import Cookies from 'js-cookie' import { useCallback, useState } from 'react' import { useTranslation } from 'react-i18next' import { rememberRegistrationSuccess } from '@/app/components/base/amplitude/registration-tracking' -import Input from '@/app/components/base/input' import { resolvePostLoginRedirect } from '@/app/signin/utils/post-login-redirect' import { validPassword } from '@/config' import { useLocale } from '@/context/i18n' @@ -125,54 +127,45 @@ const ChangePasswordForm = () => {
-
- {/* Password */} -
-
- {/* Confirm Password */} -
- -
- setConfirmPassword(e.target.value)} - placeholder={t(($) => $.confirmPasswordPlaceholder, { ns: 'login' }) || ''} - /> -
-
-
- -
-
+ + $.confirmPasswordPlaceholder, { ns: 'login' }) || ''} + /> + + +
diff --git a/web/context/i18n.spec.ts b/web/context/i18n.spec.ts index 496ac82b17e..8a440ba2073 100644 --- a/web/context/i18n.spec.ts +++ b/web/context/i18n.spec.ts @@ -3,7 +3,7 @@ import type { DocPathWithoutLang } from '@/types/doc-paths' import { renderHook } from '@testing-library/react' import { useTranslation } from '#i18n' import { getDocLanguage } from '@/i18n-config/language' -import { defaultDocBaseUrl, useDocLink } from './i18n' +import { defaultDocBaseUrl, enterpriseDocBaseUrl, useDocLink } from './i18n' const mockDeploymentEdition = vi.hoisted(() => ({ value: 'CLOUD' as 'COMMUNITY' | 'ENTERPRISE' | 'CLOUD', @@ -215,6 +215,106 @@ describe('useDocLink', () => { }) }) + describe('Enterprise documentation', () => { + beforeEach(() => { + mockDeploymentEdition.value = 'ENTERPRISE' + }) + + it('should route use documentation to the versioned enterprise documentation', () => { + const { result } = renderHook(() => useDocLink()) + + expect(result.current('/use-dify/build/workflow-chatflow')).toBe( + `${enterpriseDocBaseUrl}/en/use/build/workflow-chatflow`, + ) + }) + + it.each([ + ['/use-dify/getting-started/introduction', '/use/build/workflow-chatflow'], + ['/cli/overview', '/develop/cli/introduction'], + ['/cli/authenticate', '/develop/cli/account-users/authenticate'], + ['/cli/common-tasks', '/develop/cli/account-users/common-tasks'], + ['/cli/quick-start', '/develop/cli/account-users/quick-start'], + ] as const)('should map the renamed %s page to %s', (communityPath, enterprisePath) => { + const { result } = renderHook(() => useDocLink()) + + expect(result.current(communityPath)).toBe(`${enterpriseDocBaseUrl}/en${enterprisePath}`) + }) + + it('should convert API, plugin, and CLI documentation prefixes', () => { + const { result } = renderHook(() => useDocLink()) + + expect(result.current('/api-reference/guides/knowledge')).toBe( + `${enterpriseDocBaseUrl}/en/develop/api/guides/knowledge`, + ) + expect(result.current('/develop-plugin/getting-started/getting-started-dify-plugin')).toBe( + `${enterpriseDocBaseUrl}/en/develop/plugins/getting-started/getting-started-dify-plugin`, + ) + expect(result.current('/cli/install')).toBe(`${enterpriseDocBaseUrl}/en/develop/cli/install`) + }) + + it('should remove public product prefixes and preserve anchors', () => { + const { result } = renderHook(() => useDocLink()) + + expect(result.current('/self-host/use-dify/workspace/tools#mcp' as DocPathWithoutLang)).toBe( + `${enterpriseDocBaseUrl}/en/use/workspace/tools#mcp`, + ) + expect(result.current('/cloud/use-dify/nodes/start' as DocPathWithoutLang)).toBe( + `${enterpriseDocBaseUrl}/en/use/nodes/start`, + ) + }) + + it.each(['/cloud', '/self-host'] as const)( + 'should map the bare %s product prefix to the enterprise documentation home', + (productPrefix) => { + const { result } = renderHook(() => useDocLink()) + + expect(result.current(productPrefix as DocPathWithoutLang)).toBe( + `${enterpriseDocBaseUrl}/en/`, + ) + }, + ) + + it.each([ + '/use-dify/knowledge/knowledge-request-rate-limit', + '/cloud/use-dify/knowledge/knowledge-storage-limit', + '/cloud/use-dify/workspace/subscription-management#dify-for-education', + ] as const)('should fall back to the enterprise documentation home for %s', (communityPath) => { + const { result } = renderHook(() => useDocLink()) + + expect(result.current(communityPath as DocPathWithoutLang)).toBe( + `${enterpriseDocBaseUrl}/en/`, + ) + }) + + it('should open the enterprise documentation home when no path is provided', () => { + const { result } = renderHook(() => useDocLink()) + + expect(result.current()).toBe(`${enterpriseDocBaseUrl}/en/`) + }) + + it('should use Chinese and Japanese enterprise documentation languages', () => { + vi.mocked(useTranslation).mockReturnValue({ + i18n: { language: 'zh-Hans' }, + } as ReturnType) + vi.mocked(getDocLanguage).mockReturnValue('zh') + + const { result, rerender } = renderHook(() => useDocLink()) + expect(result.current('/use-dify/nodes/start')).toBe( + `${enterpriseDocBaseUrl}/zh/use/nodes/start`, + ) + + vi.mocked(useTranslation).mockReturnValue({ + i18n: { language: 'ja-JP' }, + } as ReturnType) + vi.mocked(getDocLanguage).mockReturnValue('ja') + rerender() + + expect(result.current('/use-dify/nodes/start')).toBe( + `${enterpriseDocBaseUrl}/ja/use/nodes/start`, + ) + }) + }) + describe('Language prefix handling', () => { it('should add /en prefix for English locale', () => { vi.mocked(useTranslation).mockReturnValue({ diff --git a/web/context/i18n.ts b/web/context/i18n.ts index 6470ed5a52d..6ab81000ca2 100644 --- a/web/context/i18n.ts +++ b/web/context/i18n.ts @@ -1,6 +1,6 @@ import type { DeploymentEdition } from '@dify/contracts/api/console/system-features/types.gen' import type { Locale } from '@/i18n-config/language' -import type { DocPathWithoutLang, DocsProduct } from '@/types/doc-paths' +import type { DocLanguage, DocPathWithoutLang, DocsProduct } from '@/types/doc-paths' import { useAtomValue } from 'jotai' import { useCallback } from 'react' import { useTranslation } from '#i18n' @@ -25,9 +25,10 @@ export const useGetPricingPageLanguage = () => { } export const defaultDocBaseUrl = 'https://docs.dify.ai' +export const enterpriseDocBaseUrl = 'https://enterprise-docs.dify.ai' export type DocPathMap = Partial> -export const getDocHomePath = () => '/home' +const getDocHomePath = () => '/home' const getCurrentDocsProduct = (deploymentEdition: DeploymentEdition): DocsProduct => { if (deploymentEdition === 'CLOUD') return 'cloud' @@ -64,18 +65,69 @@ const getProductAwarePath = (path: string, deploymentEdition: DeploymentEdition) return `/${targetProduct}${pathname}${hash}` } +const replacePathPrefix = (path: string, sourcePrefix: string, targetPrefix: string): string => { + if (path === sourcePrefix) return targetPrefix + if (!path.startsWith(`${sourcePrefix}/`)) return path + + return `${targetPrefix}${path.slice(sourcePrefix.length)}` +} + +const enterpriseDocPathOverrides: Readonly> = { + '/use-dify/getting-started/introduction': '/use/build/workflow-chatflow', + '/cli/overview': '/develop/cli/introduction', + '/cli/authenticate': '/develop/cli/account-users/authenticate', + '/cli/common-tasks': '/develop/cli/account-users/common-tasks', + '/cli/quick-start': '/develop/cli/account-users/quick-start', +} + +const unavailableEnterpriseDocPaths: ReadonlySet = new Set([ + '/use/knowledge/knowledge-request-rate-limit', + '/use/knowledge/knowledge-storage-limit', + '/use/workspace/subscription-management', +]) + +const getEnterpriseDocPath = (path: string): string => { + const { pathname, hash } = splitPathHash(path) + let targetPath = replacePathPrefix(pathname, '/cloud', '') + targetPath = replacePathPrefix(targetPath, '/self-host', '') + + if (!targetPath) return '/' + + const overriddenPath = enterpriseDocPathOverrides[targetPath] + if (overriddenPath) return `${overriddenPath}${hash}` + + targetPath = replacePathPrefix(targetPath, '/use-dify', '/use') + targetPath = replacePathPrefix(targetPath, '/api-reference', '/develop/api') + targetPath = replacePathPrefix(targetPath, '/develop-plugin', '/develop/plugins') + targetPath = replacePathPrefix(targetPath, '/cli', '/develop/cli') + + if (unavailableEnterpriseDocPaths.has(targetPath)) return '/' + + return `${targetPath}${hash}` +} + +export const getEnterpriseDocUrl = (path: string, docLanguage: DocLanguage): string => { + const targetPath = path ? getEnterpriseDocPath(path) : '/' + + return `${enterpriseDocBaseUrl}/${docLanguage}${targetPath}` +} + export const useDocLink = ( baseUrl?: string, ): ((path?: DocPathWithoutLang, pathMap?: DocPathMap) => string) => { - let baseDocUrl = baseUrl || defaultDocBaseUrl - baseDocUrl = baseDocUrl.endsWith('/') ? baseDocUrl.slice(0, -1) : baseDocUrl const locale = useLocale() const deploymentEdition = useAtomValue(deploymentEditionAtom) + const useEnterpriseDocs = deploymentEdition === 'ENTERPRISE' && !baseUrl + let baseDocUrl = baseUrl || (useEnterpriseDocs ? enterpriseDocBaseUrl : defaultDocBaseUrl) + baseDocUrl = baseDocUrl.endsWith('/') ? baseDocUrl.slice(0, -1) : baseDocUrl return useCallback( (path?: DocPathWithoutLang, pathMap?: DocPathMap): string => { const docLanguage = getDocLanguage(locale) const pathUrl = path || '' let targetPath = pathMap ? pathMap[locale] || pathUrl : pathUrl + + if (useEnterpriseDocs) return getEnterpriseDocUrl(targetPath, docLanguage) + const languagePrefix = `/${docLanguage}` if (!targetPath) { @@ -86,6 +138,6 @@ export const useDocLink = ( return `${baseDocUrl}${languagePrefix}${targetPath}` }, - [baseDocUrl, deploymentEdition, locale], + [baseDocUrl, deploymentEdition, locale, useEnterpriseDocs], ) } diff --git a/web/docs/test.md b/web/docs/test.md index 3d41ca31b24..7278d266452 100644 --- a/web/docs/test.md +++ b/web/docs/test.md @@ -39,10 +39,22 @@ Use the smallest boundary that includes the behavior owner and proves the produc - Use a real browser for layout, responsive behavior, browser-specific APIs, animation, and focus behavior that `happy-dom` cannot represent faithfully. - Follow `packages/dify-ui/README.md` for the Storybook and Vitest boundary of Dify UI primitives. -Browser Mode provides a real browser runtime for focused component tests. Using it does not by itself provide end-to-end coverage or prove integration with the running app's authentication, APIs, or persistence. - Test the behavior owner. Barrel exports, pass-through wrappers, and purely presentational children do not need separate tests when the owning feature already proves their contract. Do not repeat generic behavior already owned by Base UI, React Aria, or the browser; test Dify's integration, overrides, and known regressions. +### Browser Mode Admission + +`happy-dom` is the default choice for tests under `web/`. Use the `unit` project for pure logic, hooks, and DOM-observable component or feature behavior that does not depend on a browser's rendering engine. This split follows [Vitest test projects] and [Why Browser Mode]. + +Use the `browser` project only when the asserted contract depends on browser-owned behavior that `happy-dom` cannot represent faithfully, such as: + +- Layout geometry, CSS hit testing, responsive behavior, or pointer targeting. +- Native focus, selection, scrolling, keyboard, or pointer behavior. +- Browser APIs, observers, or animation lifecycles whose real implementation affects the result. + +Rendering UI, reducing mocks, increasing confidence, or raising coverage is not enough reason to use Browser Mode. Each `*.browser.spec.{ts,tsx}` test under `web/app/` must name the browser-owned behavior and why `happy-dom` is insufficient, exercise the smallest owner through semantic locators, and justify its additional runtime. Do not use forced interaction, fixed sleeps, private DOM or CSS assertions, or real network requests. + +Browser Mode remains a focused component or feature test and currently proves Chromium only. Use the end-to-end suite for a running application, authentication, real routing, backend APIs, persistence, or complete journeys. + ## Assert Behavior, Not Implementation - Drive state transitions through props, user interaction, URL changes, or public APIs. @@ -103,8 +115,9 @@ Mocks must preserve the public contract needed by the test. Do not mock interact ## Dify Test Setup -- Tests under `web/` run in `happy-dom` through `web/vite.config.ts` and load `web/vitest.setup.ts`. -- Tests under `packages/dify-ui/` use separate Vitest Browser Mode projects: unit specs load the package styles through `vitest.setup.ts`, while Storybook tests run stories through `@storybook/addon-vitest`. +- Following [Vite+ testing configuration], tests under `web/` use two explicit projects in `web/vite.config.ts`. Supported commands and CI select one project explicitly: `unit` runs in `happy-dom` and loads `web/vitest.setup.ts`, while `browser` runs matching `app/**/*.browser.spec.{ts,tsx}` files in Playwright Chromium and loads `web/vitest.browser.setup.ts`. Bare `vp test` runs both registered projects. +- Browser failures keep screenshots and Playwright traces under `web/.vitest-browser/`. CI uploads that directory only when failure artifacts exist; Browser Mode does not own coverage or report merging. +- Tests under `packages/dify-ui/` use two Chromium Browser Mode projects: `unit` owns focused primitive contracts and loads the package styles through `vitest.setup.ts`; `storybook` owns story render, play, and accessibility contracts through `@storybook/addon-vitest`. The names identify behavior owners, not different runtimes. - New component and feature specs should generally use a sibling `__tests__/` directory. Existing colocated utility and hook specs may follow their owning module's convention. Cross-feature integration specs belong in `web/__tests__/`. - The shared `react-i18next` mock is loaded globally. Use `createReactI18nextMock` from `web/test/i18n-mock` only when a test needs custom translations. - For `nuqs` behavior, use the helpers in `web/test/nuqs-testing.tsx` and assert URL updates. Mock `nuqs` only when URL synchronization is explicitly outside the test contract. @@ -127,19 +140,21 @@ When working across several files, order the work by dependency and verify each Run from `web/`: ```bash -# Focused spec or directory -vp test run path/to/spec-or-directory +# happy-dom; omit the path to run the full unit project +vp test run --project unit path/to/spec-or-directory -# All web tests -vp test run +# Browser Mode; omit the path to run the full browser project +vp test run --project browser path/to/spec.browser.spec.tsx -# Watch mode -vp test watch path/to/spec +# Watch mode; select browser instead for Browser Mode +vp test watch --project unit path/to/spec -# Diagnostic coverage report; not an acceptance target -vp test run --coverage path/to/spec-or-directory +# Diagnostic coverage report for the unit project; not an acceptance target +vp test run --project unit --coverage path/to/spec-or-directory ``` +Always pass `--project unit` or `--project browser`. Bare `vp test` runs both registered projects and is not the standard Web test command. + ## Review Checklist - Does each test protect a reachable product contract or meaningful regression? @@ -149,12 +164,16 @@ vp test run --coverage path/to/spec-or-directory - Is the suite deterministic, focused, and cheaper to maintain than the regression it prevents? - Would the test survive a refactor that preserves behavior? - Can the reviewer name one realistic regression and the assertion that would fail? +- For Browser Mode, is the browser-owned contract explicit, impossible to prove faithfully in `happy-dom`, and worth the additional runtime? ## References - [Vitest documentation] +- [Vitest test projects] +- [Why Browser Mode] - [Vitest Browser Mode documentation] - [Vitest Browser Mode locators] +- [Vitest Browser Mode traces] - [Storybook Vitest addon] - [Testing Library guiding principles] - [React Testing Library documentation] @@ -166,6 +185,10 @@ vp test run --coverage path/to/spec-or-directory [Testing Library guiding principles]: https://testing-library.com/docs/guiding-principles [Testing Library query guidance]: https://testing-library.com/docs/queries/about [Testing Library user-event guidance]: https://testing-library.com/docs/user-event/intro -[Vitest Browser Mode documentation]: https://vitest.dev/guide/browser -[Vitest Browser Mode locators]: https://vitest.dev/api/browser/locators -[Vitest documentation]: https://vitest.dev/guide +[Vite+ testing configuration]: https://viteplus.dev/guide/test +[Vitest Browser Mode documentation]: https://v4.vitest.dev/guide/browser +[Vitest Browser Mode locators]: https://v4.vitest.dev/api/browser/locators +[Vitest Browser Mode traces]: https://v4.vitest.dev/guide/browser/trace-view +[Vitest documentation]: https://v4.vitest.dev/guide +[Vitest test projects]: https://v4.vitest.dev/guide/projects +[Why Browser Mode]: https://v4.vitest.dev/guide/browser/why diff --git a/web/features/agent-v2/agent-detail/__tests__/navigation.spec.tsx b/web/features/agent-v2/agent-detail/__tests__/navigation.spec.tsx index 86deb56efa4..dd6c315d597 100644 --- a/web/features/agent-v2/agent-detail/__tests__/navigation.spec.tsx +++ b/web/features/agent-v2/agent-detail/__tests__/navigation.spec.tsx @@ -1,13 +1,15 @@ import type { AgentAppDetailWithSite } from '@dify/contracts/api/console/agent/types.gen' import { QueryClient, QueryClientProvider } from '@tanstack/react-query' -import { render, screen } from '@testing-library/react' +import { render, screen, waitFor, within } from '@testing-library/react' import userEvent from '@testing-library/user-event' import { AgentDetailSection, AgentDetailTop } from '../navigation' const mocks = vi.hoisted(() => ({ + deleteAgent: vi.fn(), exportAppDsl: vi.fn(), pathname: '/agents/agent-1/configure', queryData: undefined as AgentAppDetailWithSite | undefined, + replace: vi.fn(), })) vi.mock('@/app/components/app/use-export-app-dsl', () => ({ @@ -33,6 +35,7 @@ vi.mock('@/next/navigation', () => ({ usePathname: () => mocks.pathname, useRouter: () => ({ back: vi.fn(), + replace: mocks.replace, }), })) @@ -64,7 +67,7 @@ vi.mock('@/service/client', () => ({ }, delete: { mutationOptions: () => ({ - mutationFn: vi.fn(), + mutationFn: mocks.deleteAgent, }), }, put: { @@ -106,6 +109,7 @@ function renderAgentDetailSection(expand = true) { describe('AgentDetailSection', () => { beforeEach(() => { vi.clearAllMocks() + mocks.deleteAgent.mockResolvedValue({}) mocks.exportAppDsl.mockResolvedValue(undefined) mocks.pathname = '/agents/agent-1/configure' mocks.queryData = createAgent() @@ -163,6 +167,48 @@ describe('AgentDetailSection', () => { }) }) + it('returns to the roster after deleting the current agent', async () => { + const user = userEvent.setup() + renderAgentDetailSection() + + await user.click(screen.getByRole('button', { name: /agentV2\.roster\.moreActions/ })) + await user.click(screen.getByRole('menuitem', { name: 'common.operation.delete' })) + + const dialog = await screen.findByRole('alertdialog', { + name: /agentV2\.roster\.deleteDialog\.title/, + }) + await user.click(within(dialog).getByRole('button', { name: 'common.operation.delete' })) + + await waitFor(() => { + expect(mocks.replace).toHaveBeenCalledWith('/agents') + }) + expect(mocks.deleteAgent.mock.calls[0]?.[0]).toEqual({ + params: { + agent_id: 'agent-1', + }, + }) + }) + + it('keeps the current agent open when deletion fails', async () => { + const user = userEvent.setup() + mocks.deleteAgent.mockRejectedValue(new Error('Delete failed')) + renderAgentDetailSection() + + await user.click(screen.getByRole('button', { name: /agentV2\.roster\.moreActions/ })) + await user.click(screen.getByRole('menuitem', { name: 'common.operation.delete' })) + + const dialog = await screen.findByRole('alertdialog', { + name: /agentV2\.roster\.deleteDialog\.title/, + }) + await user.click(within(dialog).getByRole('button', { name: 'common.operation.delete' })) + + await waitFor(() => { + expect(mocks.deleteAgent).toHaveBeenCalled() + }) + expect(mocks.replace).not.toHaveBeenCalled() + expect(dialog).toBeInTheDocument() + }) + it('does not render more actions in collapsed sidebar mode', () => { renderAgentDetailSection(false) diff --git a/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/__tests__/dialog.spec.tsx b/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/__tests__/dialog.spec.tsx index 802ec3bb1bb..977e75e85eb 100644 --- a/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/__tests__/dialog.spec.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/__tests__/dialog.spec.tsx @@ -10,6 +10,10 @@ vi.mock('@langgenius/dify-ui/toast', () => ({ }, })) +vi.mock('@/context/i18n', () => ({ + useDocLink: () => () => 'https://docs.example.com', +})) + type CliToolDialogProps = Parameters[0] function renderCliToolDialog(props?: Partial) { @@ -129,6 +133,16 @@ describe('CliToolDialog', () => { }) describe('Actions', () => { + it('should link to documentation through the shared documentation URL', () => { + renderCliToolDialog() + + expect( + screen.getByRole('link', { + name: /agentV2\.agentDetail\.configure\.tools\.cliDialog\.learnMore/, + }), + ).toHaveAttribute('href', 'https://docs.example.com') + }) + it('should keep the form open when the backdrop is clicked', async () => { const user = userEvent.setup() const { onOpenChange } = renderCliToolDialog() diff --git a/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/dialog.tsx b/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/dialog.tsx index becb0eaa71a..da28d6ae7be 100644 --- a/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/dialog.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/orchestrate/tools/cli-tool/dialog.tsx @@ -19,6 +19,7 @@ import { Input } from '@langgenius/dify-ui/input' import { toast } from '@langgenius/dify-ui/toast' import { useCallback, useState } from 'react' import { useTranslation } from 'react-i18next' +import { useDocLink } from '@/context/i18n' import { EnvVariablesTable } from '../../advanced/env' type CliToolFormValues = { @@ -53,6 +54,7 @@ export function CliToolDialog({ }) { const { t } = useTranslation('agentV2') const { t: tCommon } = useTranslation('common') + const docLink = useDocLink() const [installCommand, setInstallCommand] = useState(tool?.installCommand ?? '') const [toolName, setToolName] = useState(tool?.name ?? '') const [envVariables, setEnvVariables] = useState(() => @@ -236,7 +238,7 @@ export function CliToolDialog({
router.replace('/agents')} /> ) diff --git a/web/features/agent-v2/roster/components/__tests__/delete-agent-dialog.spec.tsx b/web/features/agent-v2/roster/components/__tests__/delete-agent-dialog.spec.tsx new file mode 100644 index 00000000000..283ef38f0a7 --- /dev/null +++ b/web/features/agent-v2/roster/components/__tests__/delete-agent-dialog.spec.tsx @@ -0,0 +1,51 @@ +import { render, screen } from '@testing-library/react' +import { DeleteAgentDialog } from '../delete-agent-dialog' + +const mutationMock = vi.hoisted(() => ({ + isPending: false, + mutate: vi.fn(), +})) + +vi.mock('@tanstack/react-query', () => ({ + useMutation: () => mutationMock, +})) + +vi.mock('@/service/client', () => ({ + consoleQuery: { + agent: { + byAgentId: { + delete: { + mutationOptions: vi.fn(() => ({})), + }, + }, + }, + }, +})) + +vi.mock('react-i18next', async () => { + const { createReactI18nextMock } = await import('@/test/i18n-mock') + const { default: agentV2 } = await import('@/i18n/en-US/agent-v-2.json') + return createReactI18nextMock({ + 'roster.deleteDialog.description': agentV2['roster.deleteDialog.description'], + 'roster.deleteDialog.title': agentV2['roster.deleteDialog.title'], + }) +}) + +describe('DeleteAgentDialog', () => { + it('identifies the deleted agent and explains the irreversible impact', () => { + render( + , + ) + + expect( + screen.getByText( + 'This permanently deletes Research Agent. Its web app, API access, and workflows that use it stop working immediately.', + ), + ).toBeInTheDocument() + }) +}) diff --git a/web/features/agent-v2/roster/components/delete-agent-dialog.tsx b/web/features/agent-v2/roster/components/delete-agent-dialog.tsx index 3763d61dd81..7984d2c8b5a 100644 --- a/web/features/agent-v2/roster/components/delete-agent-dialog.tsx +++ b/web/features/agent-v2/roster/components/delete-agent-dialog.tsx @@ -19,6 +19,7 @@ type DeleteAgentDialogProps = { agentName: string open: boolean onOpenChange: (open: boolean) => void + onDeleted?: () => void } export function DeleteAgentDialog({ @@ -26,6 +27,7 @@ export function DeleteAgentDialog({ agentName, open, onOpenChange, + onDeleted, }: DeleteAgentDialogProps) { const { t } = useTranslation('agentV2') const { t: tCommon } = useTranslation('common') @@ -44,6 +46,7 @@ export function DeleteAgentDialog({ onSuccess: () => { toast.success(t(($) => $['roster.deleteSuccess'])) onOpenChange(false) + onDeleted?.() }, onError: () => { toast.error(t(($) => $['roster.deleteFailed'])) @@ -59,7 +62,7 @@ export function DeleteAgentDialog({ {t(($) => $['roster.deleteDialog.title'], { name: agentName })} - {t(($) => $['roster.deleteDialog.description'])} + {t(($) => $['roster.deleteDialog.description'], { name: agentName })} diff --git a/web/features/home/home-content/home-content.tsx b/web/features/home/home-content/home-content.tsx index 449bd77a9f1..e3e28e3ff47 100644 --- a/web/features/home/home-content/home-content.tsx +++ b/web/features/home/home-content/home-content.tsx @@ -349,7 +349,6 @@ export function HomeContent() { input: { params: { app_id: appId } }, }), ) - if (!appDetail) throw new Error('Recommended app not found') const { export_data, mode } = appDetail currentCreateAppModeRef.current = mode diff --git a/web/features/new-rag/__tests__/new-knowledge-list.spec.tsx b/web/features/new-rag/__tests__/new-knowledge-list.spec.tsx index 5e58c4c5ab5..e2be90f39dd 100644 --- a/web/features/new-rag/__tests__/new-knowledge-list.spec.tsx +++ b/web/features/new-rag/__tests__/new-knowledge-list.spec.tsx @@ -333,6 +333,11 @@ vi.mock('../upgrade/knowledge-upgrade-card', () => ({ ), })) +vi.mock('@/context/i18n', () => ({ + useDocLink: () => (path?: string) => `https://docs.example.com${path ?? ''}`, + useLocale: () => 'en-US', +})) + vi.mock('@/service/client', () => ({ consoleQuery: { workspaces: { @@ -916,6 +921,17 @@ describe('NewKnowledgeList', () => { ) }) + it('links the guide through the shared documentation URL', () => { + setResolvedPage() + + renderWithNuqs() + + expect(screen.getByRole('link', { name: 'dataset.newKnowledge.learnMore' })).toHaveAttribute( + 'href', + 'https://docs.example.com/use-dify/knowledge/readme', + ) + }) + it('links real knowledge spaces to the new detail shell', () => { setResolvedPage([ { diff --git a/web/features/new-rag/components/knowledge-view-switcher.tsx b/web/features/new-rag/components/knowledge-view-switcher.tsx index 9f1f07bf777..7325cae9125 100644 --- a/web/features/new-rag/components/knowledge-view-switcher.tsx +++ b/web/features/new-rag/components/knowledge-view-switcher.tsx @@ -10,6 +10,7 @@ import { SegmentedControl, SegmentedControlItem } from '@langgenius/dify-ui/segm import { useIsClient } from 'foxact/use-is-client' import { useState } from 'react' import { useTranslation } from 'react-i18next' +import { useDocLink } from '@/context/i18n' import { useNewKnowledgeGuideDismissedValue, useSetNewKnowledgeGuideDismissed, @@ -22,6 +23,7 @@ export type KnowledgeViewSwitcherProps = { export function KnowledgeViewSwitcher({ value, onChange }: KnowledgeViewSwitcherProps) { const { t } = useTranslation('dataset') + const docLink = useDocLink() const guideDismissed = useNewKnowledgeGuideDismissedValue() const setGuideDismissed = useSetNewKnowledgeGuideDismissed() const isClient = useIsClient() @@ -86,7 +88,7 @@ export function KnowledgeViewSwitcher({ value, onChange }: KnowledgeViewSwitcher
0; i--) { + if (pattern[i] !== '/') continue + let backslashCount = 0 + for (let j = i - 1; j >= 0 && pattern[j] === '\\'; j--) backslashCount++ + if (backslashCount % 2 === 0) { + closingSlashIndex = i + break + } + } + if (closingSlashIndex <= 1) return null + const source = pattern.slice(1, closingSlashIndex) + const flags = pattern.slice(closingSlashIndex + 1) + if (!REGEX_FLAGS_PATTERN.test(flags)) return null + try { + return new RegExp(source, flags) + } catch { + return null + } +} + +function createRegex(pattern, optionName) { + if (pattern.startsWith('/') && pattern.lastIndexOf('/') > 0) { + const literalRegex = parseRegexPattern(pattern) + if (literalRegex) return literalRegex + warnOnce(`[prefer-tailwind-icons] Invalid regex literal in "${optionName}": ${pattern}`) + return null + } + try { + return new RegExp(pattern) + } catch { + warnOnce(`[prefer-tailwind-icons] Invalid regex in "${optionName}": ${pattern}`) + return null + } +} + +function hasRegexMatch(value, regex) { + regex.lastIndex = 0 + return regex.test(value) +} + +function normalizeSegment(value) { + return value + .replaceAll('/', '-') + .replaceAll('_', '-') + .replace(NORMALIZE_SEGMENT_SPACES_REGEX, '') + .replace(REPEATED_DASH_REGEX, '-') + .replace(EDGE_DASH_REGEX, '') + .toLowerCase() +} + +function getIconClass(importName, source, config, globalPrefix) { + const prefix = config.prefix ?? globalPrefix + config.sourceRegex.lastIndex = 0 + config.nameRegex.lastIndex = 0 + const sourceMatch = source.match(config.sourceRegex) + const nameMatch = importName.match(config.nameRegex) + const getGroup = (...keys) => { + for (const key of keys) { + const fromName = nameMatch?.groups?.[key] + if (fromName) return fromName + const fromSource = sourceMatch?.groups?.[key] + if (fromSource) return fromSource + } + return '' + } + const iconSetPart = normalizeSegment(getGroup('set', 'iconSet')) + const iconNamePart = + camelToKebab(getGroup('name', 'icon') || importName) || camelToKebab(importName) + const variantPart = normalizeSegment(getGroup('variant')) + return [prefix, iconSetPart, iconNamePart, variantPart] + .filter(Boolean) + .join('-') + .replace(REPEATED_DASH_REGEX, '-') +} + +function normalizeLibraryConfig(config) { + const sourceRegex = createRegex(config.source, 'libraries[].source') + if (!sourceRegex) return null + const nameRegex = createRegex(config.name ?? '.*', 'libraries[].name') + if (!nameRegex) return null + return { + sourceRegex, + nameRegex, + prefix: config.prefix, + } +} + +function normalizeLibraryConfigs(configs) { + const resolved = [] + for (const config of configs) { + const normalized = normalizeLibraryConfig(config) + if (normalized) resolved.push(normalized) + } + return resolved +} + +function isNamedImportSpecifier(specifier) { + return ( + specifier.type === 'ImportSpecifier' && + specifier.imported.type === 'Identifier' && + specifier.local.type === 'Identifier' + ) +} + +function isJsxAttributeNamed(attribute, name) { + return ( + attribute.type === 'JSXAttribute' && + attribute.name.type === 'JSXIdentifier' && + attribute.name.name === name + ) +} + +function getNumericJsxAttributeValue(attribute) { + if (!attribute.value) return null + if (attribute.value.type === 'Literal' && typeof attribute.value.value === 'number') + return attribute.value.value + if ( + attribute.value.type === 'JSXExpressionContainer' && + attribute.value.expression.type === 'Literal' && + typeof attribute.value.expression.value === 'number' + ) { + return attribute.value.expression.value + } + return null +} + +function getClassNameValueText(classNames, classNameAttribute, sourceCode) { + if (!classNameAttribute?.value) return `{${JSON.stringify(classNames)}}` + if ( + classNameAttribute.value.type === 'Literal' && + typeof classNameAttribute.value.value === 'string' + ) { + const merged = `${classNames} ${classNameAttribute.value.value}`.trim() + return `{${JSON.stringify(merged)}}` + } + if (classNameAttribute.value.type === 'JSXExpressionContainer') { + const expression = classNameAttribute.value.expression + if (expression.type === 'JSXEmptyExpression') return `{${JSON.stringify(classNames)}}` + if ( + expression.type === 'CallExpression' && + expression.callee.type === 'Identifier' && + expression.callee.name === 'cn' + ) { + const existingArguments = expression.arguments.map((argument) => sourceCode.getText(argument)) + const argumentsText = [JSON.stringify(classNames), ...existingArguments].join(', ') + return `{cn(${argumentsText})}` + } + const expressionText = sourceCode.getText(expression) + const escapedClassNames = classNames + .replaceAll('\\', '\\\\') + .replaceAll('`', '\\`') + .replaceAll('${', '\\${') + return `{\`${escapedClassNames} \${${expressionText}}\`}` + } + return null +} + +function hasRuntimeReference(sourceCode, specifier) { + try { + const variable = sourceCode.getDeclaredVariables(specifier)[0] + if (!variable) return false + return variable.references.some((reference) => { + if (reference.identifier === specifier.local) return false + if (typeof reference.isTypeReference === 'boolean') return !reference.isTypeReference + if (typeof reference.isValueReference === 'boolean') return reference.isValueReference + return true + }) + } catch { + return false + } +} + +/** @type {import('eslint').Rule.RuleModule} */ +export default { + meta: { + type: 'suggestion', + hasSuggestions: true, + docs: { + description: 'Prefer Tailwind CSS icon classes over icon library components', + }, + schema: [ + { + type: 'object', + properties: { + libraries: { + type: 'array', + items: { + type: 'object', + properties: { + source: { type: 'string' }, + name: { type: 'string' }, + prefix: { type: 'string' }, + }, + required: ['source'], + additionalProperties: false, + }, + }, + prefix: { + type: 'string', + description: 'Global class prefix added before generated icon classes', + }, + propMappings: { + type: 'object', + additionalProperties: { type: 'string' }, + description: 'Maps component props to Tailwind class prefixes', + }, + }, + additionalProperties: false, + }, + ], + messages: { + preferTailwindIcon: + 'Prefer using Tailwind CSS icon class "{{iconClass}}" over "{{componentName}}" from "{{source}}"', + preferTailwindIconImport: + 'Icon "{{importedName}}" from "{{source}}" can be replaced with Tailwind CSS class "{{iconClass}}"', + }, + }, + create(context) { + const [options = {}] = context.options + const resolvedConfigs = normalizeLibraryConfigs(options.libraries ?? []) + if (resolvedConfigs.length === 0) return {} + + const globalPrefix = options.prefix ?? '' + const propMappings = options.propMappings ?? {} + const iconImports = new Map() + const sourceCode = context.sourceCode + + return { + ImportDeclaration(node) { + if (node.importKind === 'type' || typeof node.source.value !== 'string') return + const source = node.source.value + const matchedConfig = resolvedConfigs.find((config) => + hasRegexMatch(source, config.sourceRegex), + ) + if (!matchedConfig) return + + for (const specifier of node.specifiers) { + if (!isNamedImportSpecifier(specifier) || specifier.importKind === 'type') continue + const importedName = specifier.imported.name + if (!hasRegexMatch(importedName, matchedConfig.nameRegex)) continue + const localName = specifier.local.name + iconImports.set(localName, { + node: specifier, + importedName, + localName, + config: matchedConfig, + source, + used: false, + }) + } + }, + JSXOpeningElement(node) { + if (node.name.type !== 'JSXIdentifier') return + const iconInfo = iconImports.get(node.name.name) + if (!iconInfo) return + + iconInfo.used = true + const iconClass = getIconClass( + iconInfo.importedName, + iconInfo.source, + iconInfo.config, + globalPrefix, + ) + const classNameAttribute = node.attributes.find((attribute) => + isJsxAttributeNamed(attribute, 'className'), + ) + const mappedClasses = [] + const consumedMappedAttributes = new Set() + for (const [propName, classPrefix] of Object.entries(propMappings)) { + const mappedAttribute = node.attributes.find((attribute) => + isJsxAttributeNamed(attribute, propName), + ) + if (!mappedAttribute) continue + const pixelValue = getNumericJsxAttributeValue(mappedAttribute) + if (pixelValue === null) continue + mappedClasses.push(pixelToClass(pixelValue, classPrefix)) + consumedMappedAttributes.add(mappedAttribute) + } + + const classesToAdd = [iconClass, ...mappedClasses].filter(Boolean).join(' ') + const classValue = getClassNameValueText(classesToAdd, classNameAttribute, sourceCode) + if (node.parent.type !== 'JSXElement') return + + context.report({ + node, + messageId: 'preferTailwindIcon', + data: { + iconClass, + componentName: iconInfo.localName, + source: iconInfo.source, + }, + ...(classValue + ? { + suggest: [ + { + messageId: 'preferTailwindIcon', + data: { + iconClass, + componentName: iconInfo.localName, + source: iconInfo.source, + }, + fix(fixer) { + const otherAttributes = node.attributes + .filter((attribute) => { + if (attribute === classNameAttribute) return false + if (attribute.type !== 'JSXAttribute') return true + return !consumedMappedAttributes.has(attribute) + }) + .map((attribute) => sourceCode.getText(attribute)) + .join(' ') + const attrsText = otherAttributes + ? `className=${classValue} ${otherAttributes}` + : `className=${classValue}` + if (node.selfClosing) + return fixer.replaceText(node.parent, ``) + const fixes = [fixer.replaceText(node, ``)] + if (node.parent.closingElement) + fixes.push(fixer.replaceText(node.parent.closingElement, '')) + return fixes + }, + }, + ], + } + : {}), + }) + }, + 'Program:exit': () => { + for (const iconInfo of iconImports.values()) { + if (iconInfo.used || !hasRuntimeReference(sourceCode, iconInfo.node)) continue + const iconClass = getIconClass( + iconInfo.importedName, + iconInfo.source, + iconInfo.config, + globalPrefix, + ) + context.report({ + node: iconInfo.node, + messageId: 'preferTailwindIconImport', + data: { + importedName: iconInfo.importedName, + source: iconInfo.source, + iconClass, + }, + }) + } + }, + } + }, +} diff --git a/web/service/__tests__/use-snippets.spec.tsx b/web/service/__tests__/use-snippets.spec.tsx new file mode 100644 index 00000000000..e254874af9b --- /dev/null +++ b/web/service/__tests__/use-snippets.spec.tsx @@ -0,0 +1,81 @@ +import type { ReactNode } from 'react' +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { act, renderHook } from '@testing-library/react' +import { beforeEach, describe, expect, it, vi } from 'vite-plus/test' +import { useExportSnippetMutation } from '../use-snippets' + +const mockExportSnippet = vi.hoisted(() => vi.fn()) + +vi.mock('@/service/client', () => ({ + consoleClient: { + workspaces: { + current: { + customizedSnippets: { + bySnippetId: { + export: { + get: mockExportSnippet, + }, + }, + }, + }, + }, + }, + consoleQuery: { + snippets: { + key: vi.fn(() => ['snippets']), + }, + workspaces: { + current: { + customizedSnippets: { + key: vi.fn(() => ['customized-snippets']), + bySnippetId: { + export: { + get: { + mutationKey: vi.fn(() => ['customized-snippets', 'export']), + }, + }, + }, + }, + }, + }, + }, +})) + +const createWrapper = () => { + const queryClient = new QueryClient({ + defaultOptions: { + mutations: { retry: false }, + queries: { retry: false }, + }, + }) + + return function Wrapper({ children }: { children: ReactNode }) { + return {children} + } +} + +describe('useExportSnippetMutation', () => { + beforeEach(() => { + vi.clearAllMocks() + mockExportSnippet.mockResolvedValue('kind: snippet') + }) + + it('exports the requested historical workflow version', async () => { + const { result } = renderHook(() => useExportSnippetMutation(), { wrapper: createWrapper() }) + + await act(async () => { + await result.current.mutateAsync({ + snippetId: 'snippet-1', + workflowId: 'workflow-1', + }) + }) + + expect(mockExportSnippet).toHaveBeenCalledWith({ + params: { snippet_id: 'snippet-1' }, + query: { + include_secret: 'false', + workflow_id: 'workflow-1', + }, + }) + }) +}) diff --git a/web/service/client.spec.ts b/web/service/client.spec.ts index 139bfd4c6e4..6f400120fa7 100644 --- a/web/service/client.spec.ts +++ b/web/service/client.spec.ts @@ -1503,11 +1503,36 @@ describe('consoleQuery agent mutation defaults', () => { expect(queryClient.getQueryData(composerQueryKey)).toEqual(savedComposerState) }) - it('should invalidate invite option lists after deleting an agent', async () => { + it('should clear deleted agent queries and invalidate invite option lists', async () => { const consoleQuery = await loadConsoleQuery() const queryClient = new QueryClient() const invalidateQueries = vi.spyOn(queryClient, 'invalidateQueries') const deletedAgent = createAgent() + const otherAgent = createAgent({ id: 'agent-2' }) + const deletedAgentDetailQueryKey = consoleQuery.agent.byAgentId.get.queryKey({ + input: { + params: { + agent_id: deletedAgent.id, + }, + }, + }) + const deletedAgentComposerQueryKey = consoleQuery.agent.byAgentId.composer.get.queryKey({ + input: { + params: { + agent_id: deletedAgent.id, + }, + }, + }) + const otherAgentDetailQueryKey = consoleQuery.agent.byAgentId.get.queryKey({ + input: { + params: { + agent_id: otherAgent.id, + }, + }, + }) + queryClient.setQueryData(deletedAgentDetailQueryKey, deletedAgent) + queryClient.setQueryData(deletedAgentComposerQueryKey, createComposerState()) + queryClient.setQueryData(otherAgentDetailQueryKey, otherAgent) const mutationOptions = consoleQuery.agent.byAgentId.delete.mutationOptions() await mutationOptions.onSuccess?.( @@ -1524,6 +1549,9 @@ describe('consoleQuery agent mutation defaults', () => { expect(invalidateQueries).toHaveBeenCalledWith({ queryKey: consoleQuery.agent.inviteOptions.get.key(), }) + expect(queryClient.getQueryData(deletedAgentDetailQueryKey)).toBeUndefined() + expect(queryClient.getQueryData(deletedAgentComposerQueryKey)).toBeUndefined() + expect(queryClient.getQueryData(otherAgentDetailQueryKey)).toEqual(otherAgent) }) }) diff --git a/web/service/client.ts b/web/service/client.ts index 9b40250132b..62626a817de 100644 --- a/web/service/client.ts +++ b/web/service/client.ts @@ -1072,6 +1072,15 @@ export const consoleQuery: RouterUtils = createTanstackQue delete: { mutationOptions: { onSuccess: (_data, variables, _onMutateResult, context) => { + context.client.removeQueries({ + queryKey: consoleQuery.agent.byAgentId.key({ + input: { + params: { + agent_id: variables.params.agent_id, + }, + }, + }), + }) context.client.setQueriesData( { queryKey: consoleQuery.agent.get.key({ type: 'query' }), diff --git a/web/service/explore.ts b/web/service/explore.ts index 92dcc93ca3d..729627fc256 100644 --- a/web/service/explore.ts +++ b/web/service/explore.ts @@ -128,7 +128,6 @@ export const fetchAppDetail = async (id: string): Promise { } export const useExportSnippetMutation = () => { - return useMutation({ - mutationFn: ({ snippetId, include = false }) => { + return useMutation({ + mutationFn: ({ snippetId, include = false, workflowId }) => { return customizedSnippetsClient.bySnippetId.export.get({ params: { snippet_id: snippetId }, - query: { include_secret: include ? 'true' : 'false' }, + query: { + include_secret: include ? 'true' : 'false', + workflow_id: workflowId, + }, }) }, }) diff --git a/web/vite.config.ts b/web/vite.config.ts index 99e18a13170..c2c76ed05e1 100644 --- a/web/vite.config.ts +++ b/web/vite.config.ts @@ -1,5 +1,6 @@ import { fileURLToPath } from 'node:url' import { configDefaults, defineConfig, lazyPlugins } from 'vite-plus' +import { playwright } from 'vite-plus/test/browser-playwright' import { createCodeInspectorPlugin, createForceInspectorClientInjectionPlugin, @@ -11,6 +12,7 @@ import { nextStaticImageTestPlugin } from './plugins/vite/next-static-image-test const projectRoot = fileURLToPath(new URL('.', import.meta.url)) const isCI = !!process.env.CI const rootClientInjectTarget = getRootClientInjectTarget(projectRoot) +const browserTestPattern = 'app/**/*.browser.spec.{ts,tsx}' export default defineConfig(({ mode }) => { const isTest = mode === 'test' @@ -22,9 +24,7 @@ export default defineConfig(({ mode }) => { plugins: lazyPlugins(async () => { const { default: react } = await import('@vitejs/plugin-react') - if (isTest) { - return [nextStaticImageTestPlugin({ projectRoot }), react()] - } + if (isTest) return [nextStaticImageTestPlugin({ projectRoot }), react()] if (isStorybook) return [react()] @@ -80,16 +80,55 @@ export default defineConfig(({ mode }) => { // Vitest config test: { - pool: 'threads', - environment: 'happy-dom', - globals: true, - setupFiles: ['./vitest.setup.ts'], - exclude: [...configDefaults.exclude, '**/*.browser.spec.{ts,tsx}'], coverage: { provider: 'v8', reporter: isCI ? ['json', 'json-summary'] : ['text', 'json', 'json-summary'], exclude: ['**/__mocks__/**'], }, + projects: [ + { + extends: true, + test: { + name: 'unit', + pool: 'threads', + environment: 'happy-dom', + globals: true, + setupFiles: ['./vitest.setup.ts'], + exclude: [...configDefaults.exclude, browserTestPattern], + }, + }, + { + extends: true, + define: { + 'process.env': '{}', + }, + plugins: lazyPlugins(async () => { + const { default: tailwindcss } = await import('@tailwindcss/vite') + return [tailwindcss()] + }), + optimizeDeps: { + include: ['vite-plus/test/browser'], + }, + test: { + name: 'browser', + globals: true, + setupFiles: ['./vitest.browser.setup.ts'], + include: [browserTestPattern], + browser: { + enabled: true, + provider: playwright(), + instances: [{ browser: 'chromium' }], + headless: true, + screenshotDirectory: './.vitest-browser/screenshots', + screenshotFailures: true, + trace: { + mode: 'retain-on-failure', + tracesDir: './.vitest-browser/traces', + }, + }, + }, + }, + ], }, } }) diff --git a/web/vitest.browser.config.ts b/web/vitest.browser.config.ts deleted file mode 100644 index df0a9361a78..00000000000 --- a/web/vitest.browser.config.ts +++ /dev/null @@ -1,28 +0,0 @@ -import tailwindcss from '@tailwindcss/vite' -import react from '@vitejs/plugin-react' -import { defineConfig } from 'vite-plus' -import { playwright } from 'vite-plus/test/browser-playwright' - -export default defineConfig({ - define: { - 'process.env': '{}', - }, - plugins: [tailwindcss(), react()], - resolve: { - tsconfigPaths: true, - }, - optimizeDeps: { - include: ['vite-plus/test/browser'], - }, - test: { - globals: true, - setupFiles: ['./vitest.browser.setup.ts'], - include: ['app/**/*.browser.spec.{ts,tsx}'], - browser: { - enabled: true, - provider: playwright(), - instances: [{ browser: 'chromium' }], - headless: true, - }, - }, -})