diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 16d273cefb2e6a..1120ed58967284 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1515,6 +1515,10 @@ class CeleryScheduleTasksConfig(BaseSettings): description="Enable recovery of persisted workflow draft upload cleanup requests", default=True, ) + WORKFLOW_DRAFT_FILE_CLEANUP_INTERVAL: PositiveInt = Field( + description="Workflow draft upload cleanup recovery interval in minutes", + default=5, + ) ENABLE_APPLICATION_TRACE_RECONCILE_TASK: bool = Field( description="Enable periodic application trace recovery and retention cleanup when Agent tracing is enabled", default=True, diff --git a/api/controllers/common/errors.py b/api/controllers/common/errors.py index 7b061be1fe1cf6..e75c096860ce92 100644 --- a/api/controllers/common/errors.py +++ b/api/controllers/common/errors.py @@ -215,14 +215,19 @@ class RequestBodyTooLargeError(BaseHTTPException): code = 413 +def reraise_http_error(error: Exception) -> None: + """Preserve framework and application HTTP errors at a transport boundary.""" + if isinstance(error, HTTPException): + raise error + + def raise_unexpected_error(error: Exception, *, message: str) -> NoReturn: """Keep HTTP failures intact and sanitize unclassified failures at the transport boundary. Framework request parsing and project HTTP errors may both reach a route's error translator. Neither should be converted to a generic server error. """ - if isinstance(error, HTTPException): - raise error + reraise_http_error(error) logging.getLogger(__name__).exception(message) raise InternalServerError() from error diff --git a/api/controllers/mcp/mcp.py b/api/controllers/mcp/mcp.py index 8e1cd8869c3bf9..d8fe68272836cb 100644 --- a/api/controllers/mcp/mcp.py +++ b/api/controllers/mcp/mcp.py @@ -344,6 +344,15 @@ def _handle_mcp_request( client_name = f"{client_info.name}@{client_info.version}" end_user = self._create_end_user(client_name, app.tenant_id, app.id, mcp_server.id, session) + generate = AppGenerateService.generate + if isinstance(mcp_request.root, mcp_types.CallToolRequest): + services = application_services() + generate = partial( + AppGenerateService.generate, + variables=services.workflow_variables, + runtime=services.workflow_runtime, + ) + return handle_mcp_request( session, app, @@ -353,9 +362,5 @@ def _handle_mcp_request( end_user, request_id, protocol_version, - generate=partial( - AppGenerateService.generate, - variables=application_services().workflow_variables, - runtime=application_services().workflow_runtime, - ), + generate=generate, ) diff --git a/api/controllers/openapi/app_workflow.py b/api/controllers/openapi/app_workflow.py index 0966e4a2d6adaa..2af23e9951839e 100644 --- a/api/controllers/openapi/app_workflow.py +++ b/api/controllers/openapi/app_workflow.py @@ -66,7 +66,6 @@ from models.workflow import Workflow, WorkflowRun from services.errors.app import IsDraftWorkflowError, WorkflowNotFoundError from services.workflow.contracts import DraftWorkflowMissingError, WorkflowRecord -from services.workflow.dsl_import import stored_secret_ids from services.workflow.graph_check import CONTAINER_NODE_TYPES, GRAPH_MODES from services.workflow.graph_diff import draft_token from services.workflow_run_service import WorkflowRunListArgs @@ -374,7 +373,8 @@ def put(self, ctx: Context, app_id: str, env_id: str, *, body: EnvVariableSetPay if body.value == encrypter.full_mask_token() and body.value_type != EnvVariableValueType.SECRET: raise SecretMaskNotSecret() [mapping] = Workflow.normalize_environment_variable_mappings([{**body.model_dump(mode="json"), "id": env_id}]) - if mapping["value"] == HIDDEN_VALUE and env_id not in stored_secret_ids(draft): + secret_ids = {variable.id for variable in draft.environment_variables if isinstance(variable, SecretVariable)} + if mapping["value"] == HIDDEN_VALUE and env_id not in secret_ids: raise SecretMaskUnknownId() try: variable = variable_factory.build_environment_variable_from_mapping(mapping) diff --git a/api/controllers/trigger/webhook.py b/api/controllers/trigger/webhook.py index b2a4018cd2a4a7..f6345c1ebcaf17 100644 --- a/api/controllers/trigger/webhook.py +++ b/api/controllers/trigger/webhook.py @@ -4,7 +4,7 @@ from flask import jsonify, request from controllers.common.access_response import webhook_not_found_response -from controllers.common.errors import RequestBodyTooLargeError +from controllers.common.errors import RequestBodyTooLargeError, reraise_http_error from controllers.trigger import bp from core.trigger.debug.event_bus import TriggerDebugEventBus from core.trigger.debug.events import WebhookDebugEvent, build_webhook_pool_key @@ -163,6 +163,7 @@ def handle_webhook_debug(webhook_id: str): return webhook_not_found_response() except WebhookBodyTooLargeError as error: raise RequestBodyTooLargeError(str(error)) from error - except Exception: + except Exception as error: + reraise_http_error(error) logger.exception("Webhook debug processing failed for %s", webhook_id) return jsonify({"error": "Internal server error", "message": "An internal error has occurred."}), 500 diff --git a/api/controllers/web/completion.py b/api/controllers/web/completion.py index b8ca78c95a8886..2c2fd983216f64 100644 --- a/api/controllers/web/completion.py +++ b/api/controllers/web/completion.py @@ -3,9 +3,10 @@ from pydantic import BaseModel, Field, field_validator from sqlalchemy.orm import Session +from werkzeug.exceptions import InternalServerError import services -from controllers.common.errors import InternalServerError, InvalidRequestError, NotFoundError +from controllers.common.errors import InvalidRequestError, NotFoundError from controllers.common.fields import GeneratedAppResponse, SimpleResultResponse from controllers.common.schema import register_response_schema_models, register_schema_models from controllers.console.app.wraps import with_session diff --git a/api/extensions/application_services/workflow.py b/api/extensions/application_services/workflow.py index ddbe0a8028459d..60b08fc57e71df 100644 --- a/api/extensions/application_services/workflow.py +++ b/api/extensions/application_services/workflow.py @@ -304,9 +304,7 @@ def build_console_workflow_service( ), apps=apps, presence=WorkflowCollaborationRepository(redis=redis), - access=WorkflowAccessGateway( - session_factory=database_client, apps=EnterpriseConsoleAppAccess(members=rbac_members) - ), + access=WorkflowAccessGateway(apps=EnterpriseConsoleAppAccess(members=rbac_members), members=rbac_members), ) diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 3ca8857903d4f3..0c540c6d3e1a84 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -220,12 +220,6 @@ def __call__(self, *args: object, **kwargs: object) -> object: "task": "tasks.workflow_debug_reservation_tasks.recover_workflow_debug_reservations", "schedule": timedelta(minutes=1), } - if dify_config.ENABLE_WORKFLOW_DRAFT_FILE_CLEANUP_TASK: - imports.append("tasks.workflow_draft_var_tasks") - beat_schedule["workflow_draft_file_cleanup"] = { - "task": "tasks.workflow_draft_var_tasks.recover_draft_variable_file_cleanup_task", - "schedule": timedelta(minutes=5), - } if dify_config.KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED is True: from services.knowledge_fs.background_contract import OPERATION_QUEUES, OPERATION_TASK @@ -288,6 +282,12 @@ def __call__(self, *args: object, **kwargs: object) -> object: "schedule": collection_schedule, "options": {"expires": interval}, } + if dify_config.ENABLE_WORKFLOW_DRAFT_FILE_CLEANUP_TASK: + imports.append("tasks.workflow_draft_var_tasks") + beat_schedule["workflow_draft_file_cleanup"] = { + "task": "tasks.workflow_draft_var_tasks.recover_draft_variable_file_cleanup_task", + "schedule": timedelta(minutes=dify_config.WORKFLOW_DRAFT_FILE_CLEANUP_INTERVAL), + } if dify_config.ENABLE_CONVERSATION_CLEANUP_TASK: imports.append("tasks.delete_conversation_task") beat_schedule["conversation_cleanup_sweeper"] = { diff --git a/api/services/workflow/access_gateway.py b/api/services/workflow/access_gateway.py index cd29a9151e35df..7721f53a3d7f95 100644 --- a/api/services/workflow/access_gateway.py +++ b/api/services/workflow/access_gateway.py @@ -1,23 +1,27 @@ """Workflow collaboration visibility and asset delivery adapters.""" -from sqlalchemy.orm import Session, sessionmaker - from configs import dify_config from graphon.file import helpers as file_helpers from machinery.context import RequestContext from services.app.access import resolve_app_access_filter from services.app.console_service import ConsoleAppAccess +from services.rbac.members import MemberService class WorkflowAccessGateway: - def __init__(self, *, session_factory: sessionmaker[Session], apps: ConsoleAppAccess) -> None: - self._sessions = session_factory + def __init__(self, *, apps: ConsoleAppAccess, members: MemberService) -> None: self._apps = apps + self._members = members def accessible_app_ids(self, context: RequestContext, maintainers: dict[str, str | None]) -> set[str]: if not dify_config.RBAC_ENABLED: return set(maintainers) - access = resolve_app_access_filter(context.active_workspace_id, context.account_id) + permissions = self._members.permissions(context.active_workspace_id, context.account_id) + access = resolve_app_access_filter( + context.active_workspace_id, + context.account_id, + permissions=permissions, + ) return { app_id for app_id, maintainer in maintainers.items() 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 d8a82a0cf9c54b..26471ac195e616 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow.py @@ -91,8 +91,9 @@ def test_sync_invalid_input_does_not_dispatch( invoke(controller.DraftWorkflowApi, controller.DraftWorkflowApi.post) assert error.value.code == status else: - with pytest.raises(InvalidArgumentError, match="Invalid JSON data"): + with pytest.raises(InvalidArgumentError, match="Invalid JSON data") as error: invoke(controller.DraftWorkflowApi, controller.DraftWorkflowApi.post) + assert error.value.code == status workflows.sync.assert_not_called() @@ -284,13 +285,14 @@ def test_convert_delegates_permissions_result(workflows: Mock) -> None: def test_update_empty_payload_does_not_dispatch(workflows: Mock) -> None: - with pytest.raises(InvalidArgumentError, match="No valid fields to update"): + with pytest.raises(InvalidArgumentError, match="No valid fields to update") as error: invoke( controller.WorkflowByIdApi, controller.WorkflowByIdApi.patch, controller.WorkflowUpdatePayload(), workflow_id="version", ) + assert error.value.code == 400 workflows.update.assert_not_called() diff --git a/api/tests/unit_tests/controllers/web/test_completion.py b/api/tests/unit_tests/controllers/web/test_completion.py index e73e1bdc7574bc..db3b80938bb433 100644 --- a/api/tests/unit_tests/controllers/web/test_completion.py +++ b/api/tests/unit_tests/controllers/web/test_completion.py @@ -9,8 +9,8 @@ import pytest from flask import Flask from sqlalchemy.orm import Session +from werkzeug.exceptions import InternalServerError -from controllers.common.errors import InternalServerError from controllers.web.completion import ChatApi, ChatStopApi, CompletionApi, CompletionStopApi from controllers.web.error import ( AgentNotPublishedError, diff --git a/api/tests/unit_tests/extensions/test_draft_variable_cleanup_schedule.py b/api/tests/unit_tests/extensions/test_draft_variable_cleanup_schedule.py new file mode 100644 index 00000000000000..0602981f96b4ea --- /dev/null +++ b/api/tests/unit_tests/extensions/test_draft_variable_cleanup_schedule.py @@ -0,0 +1,36 @@ +"""Recovery remains scheduled even when initial cleanup delivery fails.""" + +from datetime import timedelta + +import pytest +from celery import current_app + +from dify_app import DifyApp +from extensions.ext_celery import init_app +from tests.helpers.unit.config_override import apply_config_overrides + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_draft_cleanup_recovery_schedule(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None: + restore_default = current_app.set_default + restore_current = current_app.set_current + apply_config_overrides( + monkeypatch, + ENABLE_WORKFLOW_DRAFT_FILE_CLEANUP_TASK=enabled, + WORKFLOW_DRAFT_FILE_CLEANUP_INTERVAL=7, + ) + application = init_app(DifyApp(__name__)) + try: + schedules = application.conf.beat_schedule + if enabled: + assert schedules["workflow_draft_file_cleanup"] == { + "task": "tasks.workflow_draft_var_tasks.recover_draft_variable_file_cleanup_task", + "schedule": timedelta(minutes=7), + } + assert "tasks.workflow_draft_var_tasks" in application.conf.imports + else: + assert "workflow_draft_file_cleanup" not in schedules + finally: + application.close() + restore_default() + restore_current() diff --git a/api/tests/unit_tests/extensions/test_ext_celery.py b/api/tests/unit_tests/extensions/test_ext_celery.py index fb2b3a4a992d5f..7b4139a760c659 100644 --- a/api/tests/unit_tests/extensions/test_ext_celery.py +++ b/api/tests/unit_tests/extensions/test_ext_celery.py @@ -34,7 +34,7 @@ def restore_celery_app_state() -> Iterator[None]: class TestCelerySSLConfiguration: """Test suite for Celery SSL configuration.""" - def test_get_celery_broker_transport_options_includes_global_keyprefix_for_redis(self): + def test_get_celery_broker_transport_options_includes_global_keyprefix_for_redis(self) -> None: mock_config = MagicMock() mock_config.CELERY_USE_SENTINEL = False mock_config.REDIS_KEY_PREFIX = "enterprise-a" @@ -46,7 +46,7 @@ def test_get_celery_broker_transport_options_includes_global_keyprefix_for_redis assert result["global_keyprefix"] == "enterprise-a:" - def test_get_celery_broker_transport_options_omits_global_keyprefix_when_prefix_empty(self): + def test_get_celery_broker_transport_options_omits_global_keyprefix_when_prefix_empty(self) -> None: mock_config = MagicMock() mock_config.CELERY_USE_SENTINEL = False mock_config.REDIS_KEY_PREFIX = " " @@ -58,7 +58,7 @@ def test_get_celery_broker_transport_options_omits_global_keyprefix_when_prefix_ assert "global_keyprefix" not in result - def test_get_celery_broker_transport_options_keeps_sentinel_and_adds_global_keyprefix(self): + def test_get_celery_broker_transport_options_keeps_sentinel_and_adds_global_keyprefix(self) -> None: mock_config = MagicMock() mock_config.CELERY_USE_SENTINEL = True mock_config.CELERY_SENTINEL_MASTER_NAME = "mymaster" @@ -75,7 +75,7 @@ def test_get_celery_broker_transport_options_keeps_sentinel_and_adds_global_keyp assert result["sentinel_kwargs"]["password"] == "secret" assert result["global_keyprefix"] == "enterprise-a:" - def test_get_celery_ssl_options_when_ssl_disabled(self): + def test_get_celery_ssl_options_when_ssl_disabled(self) -> None: """Test SSL options when BROKER_USE_SSL is False.""" from configs import DifyConfig @@ -87,7 +87,7 @@ def test_get_celery_ssl_options_when_ssl_disabled(self): result = get_celery_ssl_options() assert result is None - def test_get_celery_ssl_options_when_broker_not_redis(self): + def test_get_celery_ssl_options_when_broker_not_redis(self) -> None: """Test SSL options when broker is not Redis.""" mock_config = MagicMock() mock_config.CELERY_BROKER_URL = "amqp://localhost:5672" @@ -98,7 +98,7 @@ def test_get_celery_ssl_options_when_broker_not_redis(self): result = get_celery_ssl_options() assert result is None - def test_get_celery_ssl_options_with_cert_none(self): + def test_get_celery_ssl_options_with_cert_none(self) -> None: """Test SSL options with CERT_NONE requirement.""" mock_config = MagicMock() mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" @@ -117,7 +117,7 @@ def test_get_celery_ssl_options_with_cert_none(self): assert result["ssl_certfile"] is None assert result["ssl_keyfile"] is None - def test_get_celery_ssl_options_with_cert_required(self): + def test_get_celery_ssl_options_with_cert_required(self) -> None: """Test SSL options with CERT_REQUIRED and certificates.""" mock_config = MagicMock() mock_config.CELERY_BROKER_URL = "rediss://localhost:6380/0" @@ -136,7 +136,7 @@ def test_get_celery_ssl_options_with_cert_required(self): assert result["ssl_certfile"] == "/path/to/client.crt" assert result["ssl_keyfile"] == "/path/to/client.key" - def test_get_celery_ssl_options_with_cert_optional(self): + def test_get_celery_ssl_options_with_cert_optional(self) -> None: """Test SSL options with CERT_OPTIONAL requirement.""" mock_config = MagicMock() mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" @@ -153,7 +153,7 @@ def test_get_celery_ssl_options_with_cert_optional(self): assert result["ssl_cert_reqs"] == ssl.CERT_OPTIONAL assert result["ssl_ca_certs"] == "/path/to/ca.crt" - def test_get_celery_ssl_options_with_invalid_cert_reqs(self): + def test_get_celery_ssl_options_with_invalid_cert_reqs(self) -> None: """Test SSL options with invalid cert requirement defaults to CERT_NONE.""" mock_config = MagicMock() mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" @@ -169,7 +169,7 @@ def test_get_celery_ssl_options_with_invalid_cert_reqs(self): assert result is not None assert result["ssl_cert_reqs"] == ssl.CERT_NONE # Should default to CERT_NONE - def test_celery_init_applies_ssl_to_broker_and_backend(self): + def test_celery_init_applies_ssl_to_broker_and_backend(self) -> None: """Test that SSL options are applied to both broker and backend when using Redis.""" mock_config = MagicMock() mock_config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1 @@ -188,6 +188,7 @@ def test_celery_init_applies_ssl_to_broker_and_backend(self): # Mock all the scheduler configs mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 mock_config.AGENT_SANDBOX_METERING_ENABLED = False + mock_config.WORKFLOW_DRAFT_FILE_CLEANUP_INTERVAL = 5 mock_config.ENABLE_WORKFLOW_DRAFT_FILE_CLEANUP_TASK = False mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = False mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 @@ -231,7 +232,7 @@ def test_celery_init_applies_ssl_to_broker_and_backend(self): assert celery_app.conf["redis_backend_use_ssl"] is not None @pytest.mark.parametrize("draft_cleanup_enabled", [True, False]) - def test_celery_init_registers_required_agent_and_conversation_tasks(self, draft_cleanup_enabled: bool): + def test_celery_init_registers_required_agent_and_conversation_tasks(self, draft_cleanup_enabled: bool) -> None: mock_config = MagicMock() mock_config.BROKER_USE_SSL = False mock_config.REDIS_KEY_PREFIX = "enterprise-a" @@ -243,10 +244,11 @@ def test_celery_init_registers_required_agent_and_conversation_tasks(self, draft mock_config.LOG_FORMAT = "%(message)s" mock_config.LOG_TZ = "UTC" mock_config.LOG_FILE = None - mock_config.CELERY_TASK_ANNOTATIONS = {} + mock_config.CELERY_TASK_ANNOTATIONS = dict[str, object]() mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 mock_config.AGENT_SANDBOX_METERING_ENABLED = False + mock_config.WORKFLOW_DRAFT_FILE_CLEANUP_INTERVAL = 5 mock_config.ENABLE_WORKFLOW_DRAFT_FILE_CLEANUP_TASK = draft_cleanup_enabled mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = True mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 diff --git a/api/tests/unit_tests/services/annotation/test_reply_service.py b/api/tests/unit_tests/services/annotation/test_reply_service.py index 5cc598e2e0e02c..14d0df35a2d94b 100644 --- a/api/tests/unit_tests/services/annotation/test_reply_service.py +++ b/api/tests/unit_tests/services/annotation/test_reply_service.py @@ -34,7 +34,13 @@ class AnnotationStore: sessions: sessionmaker[Session] -type OwnershipChange = Literal["disable", "foreign_annotation", "foreign_message", "binding"] +type OwnershipChange = Literal[ + "binding", + "disable", + "foreign_annotation", + "foreign_message", + "missing_annotation", +] @pytest.fixture @@ -191,9 +197,18 @@ def test_search_releases_connection_and_commits_hit( hit = session.scalar(select(AppAnnotationHitHistory)) assert hit is not None assert (hit.annotation_id, hit.message_id, hit.source) == (reply.id, "msg-1", source) + assert (hit.question, hit.annotation_question, hit.annotation_content, hit.score) == ( + "hi", + "question", + "answer", + 0.8, + ) -@pytest.mark.parametrize("change", ["disable", "foreign_annotation", "foreign_message", "binding"]) +@pytest.mark.parametrize( + "change", + ["disable", "foreign_annotation", "foreign_message", "missing_annotation", "binding"], +) def test_revalidates_after_remote_search( annotation_store: AnnotationStore, monkeypatch: pytest.MonkeyPatch, @@ -213,6 +228,10 @@ def change_owner() -> None: annotation = session.get(MessageAnnotation, "annotation-1") assert annotation is not None annotation.app_id = "foreign" + elif change == "missing_annotation": + annotation = session.get(MessageAnnotation, "annotation-1") + assert annotation is not None + session.delete(annotation) else: message = session.get(Message, "msg-1") assert message is not None diff --git a/api/tests/unit_tests/services/app/test_console_gateway.py b/api/tests/unit_tests/services/app/test_console_gateway.py index a54a3e3ae000a7..51cdb3ca8e6455 100644 --- a/api/tests/unit_tests/services/app/test_console_gateway.py +++ b/api/tests/unit_tests/services/app/test_console_gateway.py @@ -36,6 +36,7 @@ from models.tool_runtime_contracts import WorkflowToolQueries from models.workflow import Workflow, WorkflowType from repositories.app.export_repository import AppExportRepository +from repositories.workflow.definition_repository import WorkflowDefinitionStore from services.agent.dsl_entities import AGENT_PACKAGE_REF_KEY from services.agent.package_resource_exporter import AgentPackageResourceExporter from services.agent.roster_package_exporter import RosterAgentPackageExporter @@ -803,9 +804,9 @@ def import_app(_dsl: AppDslService, *, preserve_agent_bindings: bool, **_kwargs: def test_create_draft_saves_the_starter_graph_of_the_apps_mode( sqlite_session_factory: sessionmaker[Session], copy_source: tuple[RequestContext, str], - monkeypatch: pytest.MonkeyPatch, tool_providers: ToolProviders, workflow_queries: WorkflowToolQueries, + monkeypatch: pytest.MonkeyPatch, mode: AppMode, ) -> None: context, app_id = copy_source @@ -815,9 +816,9 @@ def test_create_draft_saves_the_starter_graph_of_the_apps_mode( app.mode = mode syncs: list[dict[str, object]] = [] monkeypatch.setattr( - console_gateway.WorkflowDefinitionStore, + WorkflowDefinitionStore, "sync_draft_workflow", - record_calls(console_gateway.WorkflowDefinitionStore.sync_draft_workflow, syncs), + record_calls(WorkflowDefinitionStore.sync_draft_workflow, syncs), ) AppLifecycleGateway( session_factory=sqlite_session_factory, 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 1edef3a132a46f..7edbab887aeb86 100644 --- a/api/tests/unit_tests/services/test_app_dsl_service.py +++ b/api/tests/unit_tests/services/test_app_dsl_service.py @@ -1,5 +1,4 @@ from collections.abc import Callable -from dataclasses import dataclass, field from io import BytesIO from typing import cast from unittest.mock import Mock, create_autospec @@ -23,7 +22,6 @@ from models.workflow import Workflow from services.agent.dsl_entities import AgentPackage from services.agent.roster_package_entities import PackageIcon -from services.agent.workflow_contracts import WorkflowAgentBindingStore from services.app_dsl_service import AppDslService, PendingData from services.app_package_service import PreparedAppPackage from services.entities.dsl_entities import AppDslOverwriteStore, AppDslOverwriteTarget, ImportStatus @@ -31,7 +29,6 @@ from services.entities.site_dsl import SiteDsl from services.errors.app import DraftChangedError, WorkflowNotFoundError from services.errors.base import NoPermissionError -from services.workflow.contracts import DraftSyncCommand from services.workflow.draft_service import WorkflowDraftService from services.workflow.graph_diff import draft_token from tests.helpers.unit.config_override import apply_config_overrides @@ -1266,42 +1263,19 @@ def test_missing_app_section_names_the_keys_that_were_present( assert not unbound_session.in_transaction() -@dataclass -class _DraftWorkflowService: - draft: Workflow | None - lookups: list[str] = field(default_factory=list) - syncs: list[DraftSyncCommand] = field(default_factory=list) - - def _draft_of(self, app_model: App) -> Workflow | None: - return self.draft if self.draft is not None and self.draft.app_id == app_model.id else None - - def get_draft_workflow(self, app_model: App, *, for_update: bool = False, **_: object) -> Workflow | None: - self.lookups.append("locked" if for_update else "read") - return self._draft_of(app_model) - - def sync(self, _context: object, _owner: object, command: DraftSyncCommand, **_: object) -> Workflow: - self.syncs.append(command) - return _workflow(graph=command.graph) - - def _workflow_draft(app: App) -> Workflow: return make_workflow(tenant_id=app.tenant_id, app_id=app.id, graph={"nodes": [], "edges": []}) -def _guarded_overwrite( - monkeypatch: pytest.MonkeyPatch, session: Session, *, draft: Workflow | None -) -> tuple[AppDslService, _DraftWorkflowService, Callable[[str], App]]: - workflow_service = _DraftWorkflowService(draft=draft) - service = AppDslService( - session=session, - drafts=cast(WorkflowDraftService[WorkflowAgentBindingStore], workflow_service), - overwrites=create_autospec(AppDslOverwriteStore, instance=True, spec_set=True), - ) - monkeypatch.setattr( - "services.app_dsl_service.WorkflowDefinitionStore.get_draft_workflow", - workflow_service.get_draft_workflow, - ) +def _guarded_overwrite(session: Session, *, draft: Workflow | None) -> tuple[AppDslService, App, Callable[[str], App]]: + service = build_app_dsl_service(session=session) app = _app(mode=AppMode.WORKFLOW) + session.add(app) + if draft is not None: + draft.tenant_id = app.tenant_id + draft.app_id = app.id + session.add(draft) + session.flush() def overwrite(draft_hash: str) -> App: return service._create_or_update_app( @@ -1311,40 +1285,33 @@ def overwrite(draft_hash: str) -> App: draft_hash=draft_hash, ) - return service, workflow_service, overwrite + return service, app, overwrite -def test_guarded_overwrite_saves_against_the_locked_draft_it_matched( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: +def test_guarded_overwrite_saves_against_the_locked_draft_it_matched(sqlite_session: Session) -> None: draft = _workflow_draft(_app(mode=AppMode.WORKFLOW)) - _, workflow_service, overwrite = _guarded_overwrite(monkeypatch, unbound_session, draft=draft) + _, app, overwrite = _guarded_overwrite(sqlite_session, draft=draft) - overwrite(draft_token(draft)) + result = overwrite(draft_token(draft)) - assert workflow_service.lookups == ["locked"] - assert len(workflow_service.syncs) == 1 - assert workflow_service.syncs[0].graph == {"nodes": [], "edges": []} + assert result is app + assert sqlite_session.scalar(select(Workflow).where(Workflow.app_id == app.id)) is not None @pytest.mark.parametrize("draft_exists", [True, False], ids=["draft-changed", "draft-gone"]) -def test_guarded_overwrite_refuses_a_draft_that_no_longer_matches( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session, draft_exists: bool -) -> None: +def test_guarded_overwrite_refuses_a_draft_that_no_longer_matches(sqlite_session: Session, draft_exists: bool) -> None: draft = _workflow_draft(_app(mode=AppMode.WORKFLOW)) if draft_exists else None - _, workflow_service, overwrite = _guarded_overwrite(monkeypatch, unbound_session, draft=draft) + _, _app_model, overwrite = _guarded_overwrite(sqlite_session, draft=draft) with pytest.raises(DraftChangedError): overwrite("exported") - assert workflow_service.syncs == [] - def test_refused_overwrite_leaves_the_app_and_cached_dependencies_alone( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: app = _app(mode=AppMode.WORKFLOW) - service, workflow_service, _ = _guarded_overwrite(monkeypatch, unbound_session, draft=_workflow_draft(app)) + service, app, _ = _guarded_overwrite(sqlite_session, draft=_workflow_draft(app)) cached: list[str] = [] monkeypatch.setattr("services.app_dsl_service.redis_client.setex", lambda key, _expiry, _value: cached.append(key)) dependency = PluginDependency( @@ -1363,7 +1330,6 @@ def test_refused_overwrite_leaves_the_app_and_cached_dependencies_alone( assert cached == [] assert app.name == "Existing app" - assert workflow_service.syncs == [] def _refuse_changed_draft(forwarded: list[str | None]) -> Callable[..., App]: @@ -1375,18 +1341,14 @@ def create_or_update_app(*, draft_hash: str | None = None, **_: object) -> App: def test_import_app_reraises_a_changed_draft_and_keeps_the_hash_for_the_confirm( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: pending_imports: dict[str, str] = {} monkeypatch.setattr( "services.app_dsl_service.redis_client.setex", lambda key, _expiry, value: pending_imports.__setitem__(key, value), ) - service = AppDslService( - session=unbound_session, - drafts=create_autospec(WorkflowDraftService, instance=True, spec_set=True), - overwrites=create_autospec(AppDslOverwriteStore, instance=True, spec_set=True), - ) + service = build_app_dsl_service(session=sqlite_session) forwarded: list[str | None] = [] monkeypatch.setattr(service, "_create_or_update_app", _refuse_changed_draft(forwarded)) @@ -1412,9 +1374,7 @@ def test_import_app_reraises_a_changed_draft_and_keeps_the_hash_for_the_confirm( assert stored.draft_hash == "exported" -def test_confirm_import_checks_the_pending_draft_hash( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: +def test_confirm_import_checks_the_pending_draft_hash(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: pending = PendingData( tenant_id="tenant-1", account_id="account-1", @@ -1427,11 +1387,7 @@ def test_confirm_import_checks_the_pending_draft_hash( monkeypatch.setattr("services.app_dsl_service.redis_client.get", lambda _key: pending.model_dump_json()) deleted: list[str] = [] monkeypatch.setattr("services.app_dsl_service.redis_client.delete", deleted.append) - service = AppDslService( - unbound_session, - drafts=create_autospec(WorkflowDraftService, instance=True, spec_set=True), - overwrites=create_autospec(AppDslOverwriteStore, instance=True, spec_set=True), - ) + service = build_app_dsl_service(sqlite_session) forwarded: list[str | None] = [] monkeypatch.setattr(service, "_create_or_update_app", _refuse_changed_draft(forwarded)) @@ -1510,11 +1466,7 @@ def test_failed_import_does_not_report_creation( sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, confirmed: bool ) -> None: data = "version: " + CURRENT_APP_DSL_VERSION + "\nkind: app\napp: {name: Broken, mode: chat}\n" - service = AppDslService( - sqlite_session, - drafts=create_autospec(WorkflowDraftService, instance=True, spec_set=True), - overwrites=create_autospec(AppDslOverwriteStore, instance=True, spec_set=True), - ) + service = build_app_dsl_service(sqlite_session) if confirmed: pending = PendingData( tenant_id="tenant-1", account_id="account-1", import_mode="yaml-content", yaml_content=data @@ -1528,15 +1480,13 @@ def test_failed_import_does_not_report_creation( def test_pending_import_does_not_report_creation_or_cache_marker( - unbound_session: Session, monkeypatch: pytest.MonkeyPatch + sqlite_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: cache = Mock() monkeypatch.setattr("services.app_dsl_service.redis_client.setex", cache) - result = AppDslService( - unbound_session, - drafts=create_autospec(WorkflowDraftService, instance=True, spec_set=True), - overwrites=create_autospec(AppDslOverwriteStore, instance=True, spec_set=True), - ).import_app(account=_account(), import_mode="yaml-content", yaml_content=_PENDING_WORKFLOW_DSL) + result = build_app_dsl_service(sqlite_session).import_app( + account=_account(), import_mode="yaml-content", yaml_content=_PENDING_WORKFLOW_DSL + ) assert result.status == ImportStatus.PENDING assert not result.app_created assert "app_created" not in cache.call_args.args[2] 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 0519ca3b475173..316ae7adec0efc 100644 --- a/api/tests/unit_tests/services/test_snippet_dsl_service.py +++ b/api/tests/unit_tests/services/test_snippet_dsl_service.py @@ -10,7 +10,7 @@ from sqlalchemy import event, select from sqlalchemy.orm import Session -from extensions.application_services.snippets import build_snippet_dsl_service, build_snippet_service +from extensions.application_services.snippets import build_snippet_dsl_service from graphon.nodes import BuiltinNodeTypes from models import Account from models.agent import WorkflowAgentBindingType, WorkflowAgentNodeBinding @@ -553,7 +553,7 @@ def forbidden_yaml_load(*_args: object, **_kwargs: object) -> NoReturn: assert snippet.created_by == account.id assert snippet.name == "Override name" assert snippet.description == "Override description" - assert build_snippet_service(session=sqlite_session).get_draft_workflow(snippet) is not None + assert service._snippets.get_draft_workflow(snippet) is not None assert commands[-1] == ("DEL", redis_key) assert redis_key not in values @@ -723,7 +723,7 @@ def retire_unowned( persisted = sqlite_session.get(CustomizedSnippet, snippet.id) assert persisted is not None assert persisted.name == "New" - workflow = build_snippet_service(session=sqlite_session).get_draft_workflow(snippet) + workflow = service._snippets.get_draft_workflow(snippet) assert workflow is draft_workflow assert workflow.graph_dict == {"nodes": [], "edges": []} assert sqlite_session.get(WorkflowAgentNodeBinding, "retired-binding") is None @@ -755,7 +755,7 @@ def test_create_or_update_snippet_creates_new_snippet_and_flushes(service: Snipp assert result.type == "group" assert not sqlite_session.in_transaction() assert sqlite_session.get(CustomizedSnippet, result.id) is result - workflow = build_snippet_service(session=sqlite_session).get_draft_workflow(result) + workflow = service._snippets.get_draft_workflow(result) assert workflow is not None assert workflow.graph_dict == {"nodes": [], "edges": []} assert workflow.created_by == "account-1" 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 008048c04a9404..53c20ff465d46e 100644 --- a/api/tests/unit_tests/services/test_snippet_generate_service.py +++ b/api/tests/unit_tests/services/test_snippet_generate_service.py @@ -93,6 +93,8 @@ def test_debug_generators_use_injected_variable_dependencies( calls: list[str] = [] def execute(generator: WorkflowAppGenerator, **_kwargs: object) -> dict[str, str]: + assert generator._runtime is generation._runtime + assert generator._execution_driver is generation._runtime.execution_driver assert generator.draft_variable_loader == workflow_variables.workflow_loader assert generator.draft_variable_saver == workflow_variables.saver_factory saver = generator._get_draft_var_saver_factory(InvokeFrom.DEBUGGER, user, tenant_id=workflow.tenant_id) @@ -406,23 +408,6 @@ def test_generate_delegates_to_workflow_generator_and_filters_stream( assert kwargs["call_depth"] == 0 -def test_ensure_start_node_for_worker_delegates(monkeypatch: pytest.MonkeyPatch): - workflow = _workflow({"nodes": [], "edges": []}) - snippet = _snippet() - ensure_calls: list[tuple[WorkflowSnapshot, CustomizedSnippet]] = [] - - def ensure_start_node(_cls, current_workflow, current_snippet): - ensure_calls.append((current_workflow, current_snippet)) - return current_workflow - - monkeypatch.setattr(SnippetGenerateService, "_ensure_start_node", classmethod(ensure_start_node)) - - result = SnippetGenerateService.ensure_start_node_for_worker(workflow, snippet) - - assert result is workflow - assert ensure_calls == [(workflow_snapshot(workflow), snippet)] - - def test_run_draft_node_delegates_to_workflow_service( *, workflow_variables: WorkflowExecutionVariables,