From ad86351cc1d18ff08c202ff9be753f916f8b3d30 Mon Sep 17 00:00:00 2001 From: zyssyz123 <916125788@qq.com> Date: Mon, 31 Aug 2026 02:26:07 +0000 Subject: [PATCH 01/21] fix(agent): reject incompatible session snapshots (#41447) --- api/controllers/console/app/completion.py | 5 ++ api/controllers/console/app/error.py | 10 +++ api/core/app/apps/agent_app/app_generator.py | 15 ++++- api/core/app/apps/agent_app/errors.py | 18 +++++ .../apps/agent_app/runtime_request_builder.py | 21 ++++++ .../base_app_generate_response_converter.py | 8 +++ api/core/app/apps/exc.py | 7 ++ .../console/agent/test_agent_controllers.py | 34 +++++++++- .../app/apps/agent_app/test_app_generator.py | 18 +++++ .../app/apps/agent_app/test_app_runner.py | 55 ++++++++++++++-- .../agent_app/test_runtime_request_builder.py | 65 +++++++++++++++++++ .../test_based_generate_task_pipeline.py | 12 ++++ 12 files changed, 261 insertions(+), 7 deletions(-) diff --git a/api/controllers/console/app/completion.py b/api/controllers/console/app/completion.py index cfa18235cd0..5f3982a65dc 100644 --- a/api/controllers/console/app/completion.py +++ b/api/controllers/console/app/completion.py @@ -16,6 +16,7 @@ from controllers.common.schema import register_response_schema_models, register_ from controllers.console import console_ns from controllers.console.agent.app_helpers import resolve_agent_runtime_app_model from controllers.console.app.error import ( + AgentSessionConfigurationChangedError, AppUnavailableError, CompletionRequestError, ConversationCompletedError, @@ -620,6 +621,10 @@ def _raise_agent_stream_error_before_response(response): if isinstance(response, _ClosableStream): response.close() message = error_payload.get("message") + if error_payload.get("code") == AgentSessionConfigurationChangedError.error_code: + raise AgentSessionConfigurationChangedError( + str(message or AgentSessionConfigurationChangedError.description) + ) raise CompletionRequestError(str(message or "Agent App chat failed.")) return _prepend_stream_chunks(buffered, chunk, iterator) diff --git a/api/controllers/console/app/error.py b/api/controllers/console/app/error.py index 1bb6fafb224..2a84336596e 100644 --- a/api/controllers/console/app/error.py +++ b/api/controllers/console/app/error.py @@ -1,3 +1,7 @@ +from core.app.apps.agent_app.errors import ( + AGENT_SESSION_CONFIGURATION_CHANGED_ERROR_CODE, + AGENT_SESSION_CONFIGURATION_CHANGED_MESSAGE, +) from libs.exception import BaseHTTPException @@ -49,6 +53,12 @@ class CompletionRequestError(BaseHTTPException): code = 400 +class AgentSessionConfigurationChangedError(BaseHTTPException): + error_code = AGENT_SESSION_CONFIGURATION_CHANGED_ERROR_CODE + description = AGENT_SESSION_CONFIGURATION_CHANGED_MESSAGE + code = 409 + + class AppMoreLikeThisDisabledError(BaseHTTPException): error_code = "app_more_like_this_disabled" description = "The 'More like this' feature is disabled. Please refresh your page." diff --git a/api/core/app/apps/agent_app/app_generator.py b/api/core/app/apps/agent_app/app_generator.py index b2457a44a04..5381cd4b1c3 100644 --- a/api/core/app/apps/agent_app/app_generator.py +++ b/api/core/app/apps/agent_app/app_generator.py @@ -31,7 +31,11 @@ from core.agent.publish_visibility import agent_has_workflow_callable_active_sna from core.app.app_config.easy_ui_based_app.model_config.converter import ModelConfigConverter from core.app.apps.agent_app.app_config_manager import AgentAppConfigManager from core.app.apps.agent_app.app_runner import AgentAppRunner -from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError +from core.app.apps.agent_app.errors import ( + AgentAppGeneratorError, + AgentAppNotPublishedError, + AgentSessionSnapshotIncompatibleError, +) from core.app.apps.agent_app.generate_response_converter import AgentAppGenerateResponseConverter from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeRequestBuilder from core.app.apps.agent_app.session_store import AgentAppWorkspaceStore @@ -531,6 +535,15 @@ class AgentAppGenerator(MessageBasedAppGenerator): ) except GenerateTaskStoppedError: pass + except AgentSessionSnapshotIncompatibleError as error: + logger.info( + "Agent App session snapshot no longer matches the current composition", + extra={ + "agent_id": application_generate_entity.agent_id, + "conversation_id": conversation_id, + }, + ) + queue_manager.publish_error(error, PublishFrom.APPLICATION_MANAGER) except Exception as e: logger.exception("Unknown Error in Agent App generate worker") queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER) diff --git a/api/core/app/apps/agent_app/errors.py b/api/core/app/apps/agent_app/errors.py index 51b4e77116a..bdcd38abfdf 100644 --- a/api/core/app/apps/agent_app/errors.py +++ b/api/core/app/apps/agent_app/errors.py @@ -1,6 +1,24 @@ +from core.app.apps.exc import AppGenerateError + +AGENT_SESSION_CONFIGURATION_CHANGED_ERROR_CODE = "agent_session_configuration_changed" +AGENT_SESSION_CONFIGURATION_CHANGED_MESSAGE = ( + "The Agent configuration changed after this conversation started. Start a new conversation to continue." +) + + class AgentAppGeneratorError(ValueError): """Raised when an Agent App turn cannot be set up.""" class AgentAppNotPublishedError(AgentAppGeneratorError): """Raised when a public Agent App runtime is requested before publish.""" + + +class AgentSessionSnapshotIncompatibleError(AppGenerateError): + """Raised when a retained session snapshot no longer matches the current composition.""" + + error_code = AGENT_SESSION_CONFIGURATION_CHANGED_ERROR_CODE + status_code = 409 + + def __init__(self) -> None: + super().__init__(AGENT_SESSION_CONFIGURATION_CHANGED_MESSAGE) diff --git a/api/core/app/apps/agent_app/runtime_request_builder.py b/api/core/app/apps/agent_app/runtime_request_builder.py index 23b93d2f9bb..70c813d0fd5 100644 --- a/api/core/app/apps/agent_app/runtime_request_builder.py +++ b/api/core/app/apps/agent_app/runtime_request_builder.py @@ -50,6 +50,8 @@ from models.agent_config_entities import AgentSoulConfig, AgentSoulToolsConfig from models.provider_ids import ModelProviderID from services.agent.prompt_mentions import expand_prompt_mentions +from .errors import AgentSessionSnapshotIncompatibleError + class AgentAppRuntimeRequestBuildError(ValueError): """Raised when Agent App state cannot be mapped to a valid run request.""" @@ -191,6 +193,7 @@ class AgentAppRuntimeRequestBuilder: metadata=metadata, ) ) + self._validate_session_snapshot_layers(request) redacted = cast(dict[str, Any], redact_for_agent_backend_log(request)) return AgentAppRuntimeRequest( request=request, @@ -199,6 +202,24 @@ class AgentAppRuntimeRequestBuilder: binding_id=context.binding_id, ) + @staticmethod + def _validate_session_snapshot_layers(request: CreateRunRequest) -> None: + """Reject stale snapshots before they reach the Agent backend. + + Draft rows are updated in place, so their IDs cannot prove that a + retained snapshot still belongs to the current composition. Agenton + requires the ordered layer names to match exactly; enforce the same + invariant at the API boundary and return a product-level error. + """ + + snapshot = request.session_snapshot + if snapshot is None: + return + snapshot_layer_names = tuple(layer.name for layer in snapshot.layers) + composition_layer_names = tuple(layer.name for layer in request.composition.layers) + if snapshot_layer_names != composition_layer_names: + raise AgentSessionSnapshotIncompatibleError() + def _build_tool_layers( self, *, diff --git a/api/core/app/apps/base_app_generate_response_converter.py b/api/core/app/apps/base_app_generate_response_converter.py index aef54cc049e..0576bd48318 100644 --- a/api/core/app/apps/base_app_generate_response_converter.py +++ b/api/core/app/apps/base_app_generate_response_converter.py @@ -7,6 +7,7 @@ from dify_agent.protocol import RunFailureType from pydantic import JsonValue from clients.agent_backend.errors import AgentBackendError, AgentBackendRunFailedError +from core.app.apps.exc import AppGenerateError from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.task_entities import AppBlockingResponse, AppStreamResponse from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError @@ -125,6 +126,13 @@ class AppGenerateResponseConverter[TBlockingResponse: AppBlockingResponse](ABC): "message": str(e), } + if isinstance(e, AppGenerateError): + return { + "code": e.error_code, + "status": e.status_code, + "message": str(e), + } + error_responses: dict[type[Exception], dict[str, JsonValue]] = { ValueError: {"code": "invalid_param", "status": 400}, ProviderTokenNotInitError: {"code": "provider_not_initialize", "status": 400}, diff --git a/api/core/app/apps/exc.py b/api/core/app/apps/exc.py index 4187118b9bc..e5cb5d31b9a 100644 --- a/api/core/app/apps/exc.py +++ b/api/core/app/apps/exc.py @@ -1,2 +1,9 @@ +class AppGenerateError(ValueError): + """Base class for application-generation errors with a stable response contract.""" + + error_code: str + status_code: int + + class GenerateTaskStoppedError(Exception): pass diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py index c2251709838..c2b43d8cb22 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py @@ -55,7 +55,7 @@ from controllers.console.agent.roster import ( from controllers.console.app import completion as completion_controller from controllers.console.app import message as message_controller from controllers.console.app.completion import AgentBuildChatFinalizeApi, AgentChatMessageApi, AgentChatMessageStopApi -from controllers.console.app.error import CompletionRequestError +from controllers.console.app.error import AgentSessionConfigurationChangedError, CompletionRequestError from controllers.console.app.message import ( AgentChatMessageListApi, AgentMessageApi, @@ -1633,6 +1633,38 @@ def test_agent_chat_stream_preflight_raises_first_error_event() -> None: assert stream.closed is True +def test_agent_chat_stream_preflight_preserves_session_configuration_error() -> None: + class ClosableStream: + def __init__(self) -> None: + self.closed = False + self._chunks = iter( + [ + "event: ping\n\n", + ( + 'data: {"event":"error","message":"Start a new conversation to continue.",' + '"code":"agent_session_configuration_changed","status":409}\n\n' + ), + ] + ) + + def __iter__(self): + return self + + def __next__(self) -> str: + return next(self._chunks) + + def close(self) -> None: + self.closed = True + + stream = ClosableStream() + with pytest.raises(AgentSessionConfigurationChangedError) as exc_info: + completion_controller._raise_agent_stream_error_before_response(stream) + assert exc_info.value.code == 409 + assert exc_info.value.error_code == "agent_session_configuration_changed" + assert "Start a new conversation" in exc_info.value.description + assert stream.closed is True + + def test_agent_chat_stream_preflight_preserves_first_normal_event() -> None: stream = iter( ["event: ping\n\n", 'data: {"event":"message","answer":"hello"}\n\n', 'data: {"event":"message_end"}\n\n'] diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py index ca0e989d9d1..f4a0c4e90f3 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py @@ -21,6 +21,7 @@ from core.app.apps.agent_app.app_generator import ( AgentAppGenerator, AgentAppGeneratorError, ) +from core.app.apps.agent_app.errors import AgentSessionSnapshotIncompatibleError from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.app.entities.queue_entities import QueueAnnotationReplyEvent @@ -427,6 +428,23 @@ class TestGenerateWorker: self._call(generator, mocker, queue_manager) assert queue_manager.publish_error.called + def test_session_configuration_change_is_published_without_unknown_error_log( + self, + generator: AgentAppGenerator, + mocker: MockerFixture, + ) -> None: + error = AgentSessionSnapshotIncompatibleError() + self._wire(generator, mocker, run_side_effect=error) + queue_manager = mocker.MagicMock() + info_log = mocker.patch(f"{MODULE}.logger.info") + exception_log = mocker.patch(f"{MODULE}.logger.exception") + + self._call(generator, mocker, queue_manager) + + queue_manager.publish_error.assert_called_once_with(error, module.PublishFrom.APPLICATION_MANAGER) + info_log.assert_called_once() + exception_log.assert_not_called() + class TestResumeAfterFormSubmission: """ENG-638: a resume turn re-sends the paused turn's original query so the diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py index a9c27c61d47..1099a252dea 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py @@ -12,7 +12,8 @@ from typing import Any, override from unittest.mock import MagicMock import pytest -from agenton.compositor import CompositorSessionSnapshot +from agenton.compositor import CompositorSessionSnapshot, LayerSessionSnapshot +from agenton.layers import LifecycleState from dify_agent.layers.ask_human import AskHumanToolResult from dify_agent.protocol import ( AgentRunUsage, @@ -52,7 +53,8 @@ from clients.agent_backend import ( ) from core.app.apps.agent_app import app_runner as app_runner_module from core.app.apps.agent_app.app_runner import AgentAppRunner -from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeRequestBuilder +from core.app.apps.agent_app.errors import AgentSessionSnapshotIncompatibleError +from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext, AgentAppRuntimeRequestBuilder from core.app.apps.agent_app.session_store import AgentAppSessionScope, StoredAgentAppSession from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom @@ -628,6 +630,34 @@ def _dify_ctx() -> DifyRunContext: ) +def _compatible_session_snapshot() -> CompositorSessionSnapshot: + request = ( + AgentAppRuntimeRequestBuilder( + dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] + ) + .build( + AgentAppRuntimeBuildContext( + dify_context=_dify_ctx(), + agent_id="agent-1", + agent_config_snapshot_id="snap-1", + agent_soul=_soul(), + conversation_id="conv-1", + user_query="hello", + idempotency_key="msg-1", + binding_id="binding-1", + backend_binding_ref="backend-binding-1", + ) + ) + .request + ) + return CompositorSessionSnapshot( + layers=[ + LayerSessionSnapshot(name=layer.name, lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}) + for layer in request.composition.layers + ] + ) + + def _runner( client: FakeAgentBackendRunClient, store: _FakeSessionStore, @@ -1314,7 +1344,7 @@ def test_repeated_tool_calls_with_placeholder_call_id_and_reused_index_create_di def test_prior_session_snapshot_is_threaded_into_request() -> None: - prior = CompositorSessionSnapshot(layers=[]) + prior = _compatible_session_snapshot() client = FakeAgentBackendRunClient() store = _FakeSessionStore(loaded=prior) qm = _FakeQueueManager() @@ -1325,8 +1355,23 @@ def test_prior_session_snapshot_is_threaded_into_request() -> None: assert client.request.session_snapshot is prior +def test_incompatible_session_snapshot_is_rejected_before_backend_invocation() -> None: + compatible = _compatible_session_snapshot() + stale = CompositorSessionSnapshot( + layers=[layer for layer in compatible.layers if layer.name != "agent_soul_prompt"] + ) + client = FakeAgentBackendRunClient() + store = _FakeSessionStore(loaded=stale) + + with pytest.raises(AgentSessionSnapshotIncompatibleError, match="Start a new conversation"): + _run(_runner(client, store), _FakeQueueManager()) + + assert client.request is None + assert store.saved == [] + + def test_debug_session_scope_can_reuse_conversation_across_config_snapshots() -> None: - prior = CompositorSessionSnapshot(layers=[]) + prior = _compatible_session_snapshot() client = FakeAgentBackendRunClient() store = _FakeSessionStore(loaded=prior) qm = _FakeQueueManager() @@ -1599,7 +1644,7 @@ def test_ask_human_pauses_turn_creates_form_and_persists_correlation() -> None: def test_submitted_form_resumes_turn_with_deferred_tool_results(monkeypatch: pytest.MonkeyPatch) -> None: # ENG-638: a turn that runs while a pending form is answered threads the # human's reply into the request as deferred_tool_results. - snapshot = CompositorSessionSnapshot(layers=[]) + snapshot = _compatible_session_snapshot() stored = StoredAgentAppSession( scope=AgentAppSessionScope( tenant_id="tenant-1", diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py index 987c15d7a39..db33a1ad03b 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py @@ -6,6 +6,8 @@ from __future__ import annotations from types import SimpleNamespace import pytest +from agenton.compositor import CompositorSessionSnapshot, LayerSessionSnapshot +from agenton.layers import LifecycleState from dify_agent.layers.config import DifyConfigSkillConfig from dify_agent.layers.dify_core_tools import DifyCoreToolConfig, DifyCoreToolsLayerConfig from dify_agent.layers.dify_plugin import DifyPluginToolConfig, DifyPluginToolsLayerConfig @@ -20,6 +22,7 @@ from clients.agent_backend import ( AgentBackendRunRequestBuilder, ) from clients.agent_backend.request_builder import DIFY_SHELL_LAYER_ID +from core.app.apps.agent_app.errors import AgentSessionSnapshotIncompatibleError from core.app.apps.agent_app.runtime_request_builder import ( AgentAppRuntimeBuildContext, AgentAppRuntimeRequestBuilder, @@ -164,6 +167,7 @@ def _ctx( *, query: str = "hello", agent_config_version_kind: str = "snapshot", + session_snapshot: CompositorSessionSnapshot | None = None, ) -> AgentAppRuntimeBuildContext: dify_context = SimpleNamespace( tenant_id="tenant-1", @@ -183,6 +187,7 @@ def _ctx( binding_id="binding-1", backend_binding_ref="binding-ref-1", agent_config_version_kind=agent_config_version_kind, # type: ignore[arg-type] + session_snapshot=session_snapshot, ) @@ -199,6 +204,15 @@ def _soul_with_model() -> AgentSoulConfig: ) +def _snapshot_for_layer_names(layer_names: list[str]) -> CompositorSessionSnapshot: + return CompositorSessionSnapshot( + layers=[ + LayerSessionSnapshot(name=name, lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}) + for name in layer_names + ] + ) + + class TestAgentAppRuntimeRequestBuilder: def test_build_maps_soul_to_run_request(self, model_context_window_calls: list[tuple[object, str, str]]): builder = AgentAppRuntimeRequestBuilder( @@ -237,6 +251,57 @@ class TestAgentAppRuntimeRequestBuilder: assert "credentials" not in result.redacted_request["composition"]["layers"][-1]["config"] assert result.metadata["conversation_id"] == "conv-1" + @pytest.mark.parametrize( + ("previous_prompt", "current_prompt"), + [("", "You are Iris."), ("You are Iris.", "")], + ) + def test_build_rejects_session_snapshot_after_layer_topology_changes( + self, + previous_prompt: str, + current_prompt: str, + ) -> None: + builder = AgentAppRuntimeRequestBuilder( + dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] + ) + previous_soul = _soul_with_model() + previous_soul.prompt.system_prompt = previous_prompt + previous_request = builder.build(_ctx(previous_soul, agent_config_version_kind="draft")).request + snapshot = _snapshot_for_layer_names([layer.name for layer in previous_request.composition.layers]) + current_soul = _soul_with_model() + current_soul.prompt.system_prompt = current_prompt + + with pytest.raises(AgentSessionSnapshotIncompatibleError) as exc_info: + builder.build( + _ctx( + current_soul, + agent_config_version_kind="draft", + session_snapshot=snapshot, + ) + ) + + assert exc_info.value.error_code == "agent_session_configuration_changed" + assert exc_info.value.status_code == 409 + assert "Start a new conversation" in str(exc_info.value) + + def test_build_reuses_session_snapshot_when_config_changes_without_changing_layers(self) -> None: + builder = AgentAppRuntimeRequestBuilder( + dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] + ) + previous_request = builder.build(_ctx(_soul_with_model(), agent_config_version_kind="draft")).request + snapshot = _snapshot_for_layer_names([layer.name for layer in previous_request.composition.layers]) + current_soul = _soul_with_model() + current_soul.prompt.system_prompt = "You are Ada." + + result = builder.build( + _ctx( + current_soul, + agent_config_version_kind="draft", + session_snapshot=snapshot, + ) + ) + + assert result.request.session_snapshot is snapshot + def test_build_wraps_agent_soul_prompt_for_build_draft(self): builder = AgentAppRuntimeRequestBuilder( dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py b/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py index 18c2fdd97a2..26b3f521a3e 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py @@ -7,6 +7,7 @@ from dify_agent.protocol import RunFailureType from sqlalchemy.orm import Session from clients.agent_backend.errors import AgentBackendRunFailedError +from core.app.apps.agent_app.errors import AgentSessionSnapshotIncompatibleError from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import QueueErrorEvent @@ -168,6 +169,17 @@ class TestBasedGenerateTaskPipeline: "message": "run limit reached (agent_run_id=run-1)", } + def test_stream_converter_preserves_agent_session_configuration_error(self): + data = AppGenerateResponseConverter._error_to_stream_response(AgentSessionSnapshotIncompatibleError()) + + assert data == { + "code": "agent_session_configuration_changed", + "status": 409, + "message": ( + "The Agent configuration changed after this conversation started. Start a new conversation to continue." + ), + } + def test_handle_output_moderation_when_flagged(self, pipeline): handler = Mock() handler.moderation_completion.return_value = ("filtered", True) From b70ad7d4240798b0b5f254f56b7b8962c29b5b8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= Date: Mon, 31 Aug 2026 02:36:32 +0000 Subject: [PATCH 02/21] fix(web): show DSL import warning details (#41502) --- web/hooks/use-import-dsl.spec.tsx | 48 +++++++++++++++++++++++++++++-- web/hooks/use-import-dsl.ts | 13 +++++++-- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/web/hooks/use-import-dsl.spec.tsx b/web/hooks/use-import-dsl.spec.tsx index 1f6e95205f9..d8aff462318 100644 --- a/web/hooks/use-import-dsl.spec.tsx +++ b/web/hooks/use-import-dsl.spec.tsx @@ -1,4 +1,4 @@ -import { act, waitFor } from '@testing-library/react' +import { act, render, screen, waitFor } from '@testing-library/react' import { DSLImportMode, DSLImportStatus } from '@/models/app' import { renderHookWithConsoleQuery } from '@/test/console/query-data' import { AppModeEnum } from '@/types/app' @@ -93,6 +93,46 @@ describe('useImportDSL', () => { mockResolveImportedAppRedirectionTarget.mockImplementation(async (target) => target) }) + it('should show response warnings when an import completes with warnings', async () => { + const completedResponse = { + id: 'import-1', + status: DSLImportStatus.COMPLETED_WITH_WARNINGS, + app_id: 'app-1', + app_mode: AppModeEnum.WORKFLOW, + permission_keys: [], + warnings: [ + { + code: 'agent_tool_authorization_required', + path: 'agent_packages.agent_1.soul.tools.dify_tools.0', + message: "Agent tool 'jina_search' requires authorization.", + details: { tool_name: 'jina_search' }, + }, + ], + } + mockImportDSL.mockResolvedValue(completedResponse) + mockHandleCheckPluginDependencies.mockResolvedValue(undefined) + + const { result } = renderHookWithConsoleQuery(() => useImportDSL()) + + await act(async () => { + await result.current.handleImportDSL( + { + mode: DSLImportMode.YAML_CONTENT, + yaml_content: 'app: demo', + }, + { skipRedirectOnSuccess: true }, + ) + }) + + expect(toastMocks.warning).toHaveBeenCalledWith('app.newApp.caution', { + description: expect.anything(), + }) + const warningDescription = toastMocks.warning.mock.calls[0]![1].description + render(<>{warningDescription}) + expect(screen.getByText("Agent tool 'jina_search' requires authorization.")).toBeInTheDocument() + expect(screen.queryByText('app.newApp.appCreateDSLWarning')).not.toBeInTheDocument() + }) + it('should complete a confirmed import that returns warnings', async () => { let resolvePluginCheck: (() => void) | undefined const pendingResponse = { @@ -163,8 +203,12 @@ describe('useImportDSL', () => { expect(onSuccess).toHaveBeenCalledWith(completedResponse) expect(onFailed).not.toHaveBeenCalled() expect(toastMocks.warning).toHaveBeenCalledWith('app.newApp.caution', { - description: 'app.newApp.appCreateDSLWarning', + description: expect.anything(), }) + const warningDescription = toastMocks.warning.mock.calls[0]![1].description + render(<>{warningDescription}) + expect(screen.getByText('Agent file was not included.')).toBeInTheDocument() + expect(screen.queryByText('app.newApp.appCreateDSLWarning')).not.toBeInTheDocument() expect(mockHandleCheckPluginDependencies).toHaveBeenCalledWith('app-1') expect(mockResolveImportedAppRedirectionTarget).toHaveBeenCalledWith({ id: 'app-1', diff --git a/web/hooks/use-import-dsl.ts b/web/hooks/use-import-dsl.ts index dda857efdb5..81e32d46a91 100644 --- a/web/hooks/use-import-dsl.ts +++ b/web/hooks/use-import-dsl.ts @@ -3,8 +3,9 @@ import type { AppIconType } from '@/types/app' import { toast } from '@langgenius/dify-ui/toast' import { useMutation, useSuspenseQuery } from '@tanstack/react-query' import { useAtomValue } from 'jotai' -import { useCallback, useRef, useState } from 'react' +import { createElement, useCallback, useRef, useState } from 'react' import { useTranslation } from 'react-i18next' +import DSLImportWarningDescription from '@/app/components/app/create-from-dsl-modal/dsl-import-warning-description' import { usePluginDependencies } from '@/app/components/workflow/plugin-dependency/hooks' import { workspacePermissionKeysAtom } from '@/context/permission-state' import { userProfileQueryOptions } from '@/features/account-profile/client' @@ -80,7 +81,10 @@ export const useImportDSL = () => { ) const description = status === DSLImportStatus.COMPLETED_WITH_WARNINGS - ? t(($) => $['newApp.appCreateDSLWarning'], { ns: 'app' }) + ? createElement(DSLImportWarningDescription, { + warnings: response.warnings, + fallback: t(($) => $['newApp.appCreateDSLWarning'], { ns: 'app' }), + }) : undefined if (status === DSLImportStatus.COMPLETED) toast.success(message) @@ -162,7 +166,10 @@ export const useImportDSL = () => { ) const description = status === DSLImportStatus.COMPLETED_WITH_WARNINGS - ? t(($) => $['newApp.appCreateDSLWarning'], { ns: 'app' }) + ? createElement(DSLImportWarningDescription, { + warnings: response.warnings, + fallback: t(($) => $['newApp.appCreateDSLWarning'], { ns: 'app' }), + }) : undefined if (status === DSLImportStatus.COMPLETED) toast.success(message) From 750810ea812dd7a4c5606e4d16cc2331eec1838a Mon Sep 17 00:00:00 2001 From: Xiyuan Chen <52963600+GareArc@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:51:19 +0000 Subject: [PATCH 03/21] fix(api): keep workspace owner consistent on ownership transfer (#41267) --- api/services/account_service.py | 70 ++++++++++-------- api/services/enterprise/rbac_service.py | 2 +- .../console/workspace/test_members.py | 2 +- .../services/test_account_service.py | 4 +- .../services/enterprise/test_rbac_service.py | 2 +- .../services/test_account_service.py | 73 ++++++++++++++++++- 6 files changed, 118 insertions(+), 35 deletions(-) diff --git a/api/services/account_service.py b/api/services/account_service.py index a60c2007912..586619abe1d 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -170,6 +170,16 @@ class AccountService: OWNER_TRANSFER_MAX_ERROR_LIMITS = 5 EMAIL_REGISTER_MAX_ERROR_LIMITS = 5 + @staticmethod + def _resolve_role_id_by_tag(tenant_id: str, account_id: str, tag: str) -> str: + options = ListOption(page_number=1, results_per_page=100) + roles = RBACService.Roles.list(tenant_id, account_id, options=options).data + for rbac_role in roles: + if rbac_role.is_builtin and rbac_role.category == "global_system_default" and rbac_role.role_tag == tag: + return str(rbac_role.id) + + raise ValueError(f"Builtin RBAC role not found for tag {tag!r} in tenant {tenant_id}") + @staticmethod def _resolve_legacy_role_id(tenant_id: str, account_id: str, role: TenantAccountRole) -> str: """Resolve a legacy workspace role to the corresponding RBAC role id. @@ -177,9 +187,6 @@ class AccountService: Looks up the builtin RBAC role whose tag matches the legacy role name (e.g. ``TenantAccountRole.ADMIN`` → builtin role with tag ``"admin"``). """ - options = ListOption(page_number=1, results_per_page=100) - roles = RBACService.Roles.list(tenant_id, account_id, options=options).data - expected_tag = { TenantAccountRole.OWNER: "owner", TenantAccountRole.ADMIN: "admin", @@ -187,15 +194,7 @@ class AccountService: TenantAccountRole.NORMAL: "normal", TenantAccountRole.DATASET_OPERATOR: "dataset_operator", }[role] - for rbac_role in roles: - if ( - rbac_role.is_builtin - and rbac_role.category == "global_system_default" - and rbac_role.role_tag == expected_tag - ): - return str(rbac_role.id) - - raise ValueError(f"Builtin RBAC role not found for {role.value} in tenant {tenant_id}") + return AccountService._resolve_role_id_by_tag(tenant_id, account_id, expected_tag) @staticmethod def get_workspace_permission_keys(tenant_id: str, account_id: str, *, session: Session) -> set[str]: @@ -1857,28 +1856,39 @@ class TenantService: raise RoleAlreadyAssignedError("The provided role is already assigned to the member.") if new_role == "owner": - # Find the current owner and change their role to 'admin' + if dify_config.RBAC_ENABLED: + old_owner_id = AccountService.get_rbac_workspace_owner_account_id( + str(tenant.id), operator.id, session=session + ) + owner_role_id = AccountService._resolve_legacy_role_id( + tenant_id=str(tenant.id), + account_id=operator.id, + role=TenantAccountRole.OWNER, + ) + no_access_role_id = AccountService._resolve_role_id_by_tag( + tenant_id=str(tenant.id), + account_id=operator.id, + tag="no_access", + ) + current_roles = RBACService.MemberRoles.get( + str(tenant.id), operator.id, old_owner_id, session=session + ).roles + remaining_role_ids = [str(r.id) for r in current_roles if str(r.id) != owner_role_id] + RBACService.MemberRoles.replace( + tenant_id=str(tenant.id), + account_id=operator.id, + member_account_id=old_owner_id, + role_ids=remaining_role_ids or [no_access_role_id], + session=session, + ) + current_owner_join = session.scalar( select(TenantAccountJoin) .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role == "owner") .limit(1) ) - if not dify_config.RBAC_ENABLED: - if current_owner_join: - current_owner_join.role = TenantAccountRole.ADMIN - elif current_owner_join: - admin_role_id = AccountService._resolve_legacy_role_id( - tenant_id=str(tenant.id), - account_id=operator.id, - role=TenantAccountRole.ADMIN, - ) - RBACService.MemberRoles.replace( - tenant_id=str(tenant.id), - account_id=operator.id, - member_account_id=str(current_owner_join.account_id), - role_ids=[admin_role_id], - session=session, - ) + if current_owner_join: + current_owner_join.role = TenantAccountRole.NORMAL # Update the role of the target member if dify_config.RBAC_ENABLED: @@ -1894,6 +1904,8 @@ class TenantService: role_ids=[resolved_role_id], session=session, ) + if new_tenant_role == TenantAccountRole.OWNER: + target_member_join.role = new_tenant_role else: target_member_join.role = new_tenant_role session.commit() diff --git a/api/services/enterprise/rbac_service.py b/api/services/enterprise/rbac_service.py index 5f206548cc5..b9a6db00e00 100644 --- a/api/services/enterprise/rbac_service.py +++ b/api/services/enterprise/rbac_service.py @@ -1846,7 +1846,7 @@ class RBACService: ) ) if current_owner_join and current_owner_join.account_id != member_account_id: - current_owner_join.role = TenantAccountRole.ADMIN + current_owner_join.role = TenantAccountRole.NORMAL target_member_join.role = tenant_role session.commit() diff --git a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_members.py b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_members.py index 1c84b70b08e..33bbdcb2f69 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_members.py +++ b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_members.py @@ -300,7 +300,7 @@ class TestOwnerTransferApiWithContainers: ) assert ( factory.get_join(db_session_with_containers, tenant=tenant, account=current_user).role - == TenantAccountRole.ADMIN + == TenantAccountRole.NORMAL ) mock_new_owner_email.assert_called_once() mock_old_owner_email.assert_called_once() diff --git a/api/tests/test_containers_integration_tests/services/test_account_service.py b/api/tests/test_containers_integration_tests/services/test_account_service.py index b8d5bcf7668..db538d78c13 100644 --- a/api/tests/test_containers_integration_tests/services/test_account_service.py +++ b/api/tests/test_containers_integration_tests/services/test_account_service.py @@ -1671,7 +1671,7 @@ class TestTenantService: def test_update_member_role_to_owner(self, db_session_with_containers: Session, mock_external_service_dependencies): """ - Test updating member role to owner (should change current owner to admin). + Test updating member role to owner (should change current owner to normal). """ fake = Faker() tenant_name = fake.company() @@ -1723,7 +1723,7 @@ class TestTenantService: .filter_by(tenant_id=tenant.id, account_id=member_account.id) .first() ) - assert owner_join.role == "admin" + assert owner_join.role == "normal" assert member_join.role == "owner" def test_update_member_role_already_assigned( diff --git a/api/tests/unit_tests/services/enterprise/test_rbac_service.py b/api/tests/unit_tests/services/enterprise/test_rbac_service.py index 93e417dbec1..432466a9975 100644 --- a/api/tests/unit_tests/services/enterprise/test_rbac_service.py +++ b/api/tests/unit_tests/services/enterprise/test_rbac_service.py @@ -997,7 +997,7 @@ class TestMemberRoles: } assert persisted_joins == { "acct-2": svc.TenantAccountRole.OWNER, - "acct-owner": svc.TenantAccountRole.ADMIN, + "acct-owner": svc.TenantAccountRole.NORMAL, } assert out.roles[0].id == "owner" diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 7a926eba6a2..a17d285a7f9 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -27,7 +27,7 @@ from services.account_service import ( RegisterService, TenantService, ) -from services.enterprise.rbac_service import MembersInRole, Paginated +from services.enterprise.rbac_service import MemberRolesResponse, MembersInRole, Paginated, RBACRole from services.errors.account import ( AccountAlreadyInTenantError, AccountEmailAlreadyInUseError, @@ -798,6 +798,14 @@ class TestTenantService: sqlite_session.add(tenant_account_join) return tenant_account_join + def _db_role_of(self, sqlite_session: Session, tenant: Tenant, account_id: str) -> str | None: + return sqlite_session.scalar( + select(TenantAccountJoin.role).where( + TenantAccountJoin.tenant_id == tenant.id, + TenantAccountJoin.account_id == account_id, + ) + ) + def test_iter_member_account_id_batches_uses_offset_limit(self, sqlite_session: Session) -> None: tenant_id = "00000000-0000-0000-0000-000000000001" account_ids = [ @@ -1332,6 +1340,69 @@ class TestTenantService: assert persisted_target_join is not None assert persisted_target_join.role == TenantAccountRole.ADMIN + @pytest.mark.parametrize( + ("outgoing_owner_role_tags", "expected_demoted_role_ids"), + [(["owner", "editor"], ["editor-role-id"]), (["owner"], ["no-access-role-id"])], + ) + def test_update_member_role_to_owner_rbac_enabled( + self, + sqlite_session: Session, + outgoing_owner_role_tags: list[str], + expected_demoted_role_ids: list[str], + config_overrides: Callable[..., None], + ) -> None: + config_overrides(RBAC_ENABLED=True) + tenant = Tenant(name="Test Workspace") + sqlite_session.add(tenant) + sqlite_session.flush() + + operator = TestAccountAssociatedDataFactory.create_account_mock(account_id="operator-1") + candidate = TestAccountAssociatedDataFactory.create_account_mock(account_id="candidate-1") + self._add_tenant_account_join(sqlite_session, tenant, operator.id, TenantAccountRole.EDITOR) + self._add_tenant_account_join(sqlite_session, tenant, candidate.id, TenantAccountRole.EDITOR) + self._add_tenant_account_join(sqlite_session, tenant, "stale-db-owner", TenantAccountRole.OWNER) + sqlite_session.commit() + + outgoing_owner_roles = MemberRolesResponse( + account_id="real-rbac-owner", + roles=[ + RBACRole(id=f"{tag}-role-id", type="workspace", name=tag, role_tag=tag) + for tag in outgoing_owner_role_tags + ], + ) + + with ( + patch( + "services.account_service.AccountService.get_workspace_permission_keys", + return_value={"workspace.role.manage"}, + ), + patch( + "services.account_service.AccountService.get_rbac_workspace_owner_account_id", + return_value="real-rbac-owner", + ), + patch( + "services.account_service.AccountService._resolve_legacy_role_id", + side_effect=lambda *, role, **_kwargs: f"{role.value}-role-id", + ), + patch( + "services.account_service.AccountService._resolve_role_id_by_tag", + return_value="no-access-role-id", + ), + patch("services.account_service.RBACService.MemberRoles.get", return_value=outgoing_owner_roles), + patch("services.account_service.RBACService.MemberRoles.replace") as mock_replace, + ): + TenantService.update_member_role(tenant, candidate, "owner", operator, session=sqlite_session) + + mock_replace.assert_any_call( + tenant_id=tenant.id, + account_id=operator.id, + member_account_id="real-rbac-owner", + role_ids=expected_demoted_role_ids, + session=sqlite_session, + ) + assert self._db_role_of(sqlite_session, tenant, "stale-db-owner") == TenantAccountRole.NORMAL + assert self._db_role_of(sqlite_session, tenant, candidate.id) == TenantAccountRole.OWNER + def test_create_owner_tenant_rbac_enabled_assigns_owner_role( self, sqlite_session: Session, From 6a00c23e492a245e85815dd8cb948407df77411b Mon Sep 17 00:00:00 2001 From: wangxiaolei Date: Mon, 31 Aug 2026 03:05:19 +0000 Subject: [PATCH 04/21] feat: use app.acl.access_point_manage to control access point page access (#41462) Co-authored-by: twwu --- api/controllers/console/agent/roster.py | 12 ++ api/services/enterprise/rbac_service.py | 3 + .../console/agent/test_agent_controllers.py | 109 +++++++++++++++++- .../[appId]/__tests__/layout-main.spec.tsx | 48 +++++++- .../(appDetailLayout)/[appId]/layout-main.tsx | 47 ++++---- .../__tests__/app-detail-section.spec.tsx | 18 ++- .../app-sidebar/app-detail-section.tsx | 16 ++- .../environment-deployment-flow.spec.tsx | 2 + .../app-publisher/__tests__/sections.spec.tsx | 41 +++++++ .../built-in-publisher/actions-section.tsx | 20 ++-- .../actions-section.tsx | 20 ++-- .../environment-deployment-flow/index.tsx | 3 + .../components/app/app-publisher/index.tsx | 8 +- .../app-publisher/publisher-content/index.tsx | 4 + .../app/deploy/__tests__/index.spec.tsx | 16 ++- .../built-in-environment-card/index.tsx | 6 +- .../app/deploy/environment-table/index.tsx | 3 + .../app/deploy/environment-table/row.tsx | 8 +- web/app/components/app/deploy/index.tsx | 13 ++- .../__tests__/access-point-icon.spec.tsx | 35 ++++++ .../app/deploy/shared/access-point-icon.tsx | 8 +- .../__tests__/index.spec.tsx | 1 + .../continue-work/__tests__/item.spec.tsx | 13 ++- web/i18n/ar-TN/permission-keys.json | 1 + web/i18n/de-DE/permission-keys.json | 1 + web/i18n/en-US/permission-keys.json | 1 + web/i18n/es-ES/permission-keys.json | 1 + web/i18n/fa-IR/permission-keys.json | 1 + web/i18n/fr-FR/permission-keys.json | 1 + web/i18n/hi-IN/permission-keys.json | 1 + web/i18n/id-ID/permission-keys.json | 1 + web/i18n/it-IT/permission-keys.json | 1 + web/i18n/ja-JP/permission-keys.json | 1 + web/i18n/ko-KR/permission-keys.json | 1 + web/i18n/lo-LA/permission-keys.json | 1 + web/i18n/nl-NL/permission-keys.json | 1 + web/i18n/pl-PL/permission-keys.json | 1 + web/i18n/pt-BR/permission-keys.json | 1 + web/i18n/ro-RO/permission-keys.json | 1 + web/i18n/ru-RU/permission-keys.json | 1 + web/i18n/sl-SI/permission-keys.json | 1 + web/i18n/th-TH/permission-keys.json | 1 + web/i18n/tr-TR/permission-keys.json | 1 + web/i18n/uk-UA/permission-keys.json | 1 + web/i18n/vi-VN/permission-keys.json | 1 + web/i18n/zh-Hans/permission-keys.json | 1 + web/i18n/zh-Hant/permission-keys.json | 1 + web/utils/app-redirection.spec.ts | 30 ++++- web/utils/app-redirection.ts | 4 +- web/utils/permission.spec.ts | 10 ++ web/utils/permission.ts | 7 ++ 51 files changed, 450 insertions(+), 79 deletions(-) create mode 100644 web/app/components/app/deploy/shared/__tests__/access-point-icon.spec.tsx diff --git a/api/controllers/console/agent/roster.py b/api/controllers/console/agent/roster.py index f95f9d8b00f..730dd7b1a57 100644 --- a/api/controllers/console/agent/roster.py +++ b/api/controllers/console/agent/roster.py @@ -7,6 +7,7 @@ from pydantic import AliasChoices, BaseModel, Field, field_validator from sqlalchemy import func, or_, select from sqlalchemy.orm import Session +from configs import dify_config from controllers.common.schema import ( query_params_from_model, query_params_from_request, @@ -78,9 +79,11 @@ from services.agent.observability_service import ( ) from services.agent.roster_service import AgentRosterService from services.app_service import AgentAppPublicationCounts, AppListParams, AppService, CreateAppParams +from services.enterprise import rbac_service as enterprise_rbac_service from services.enterprise.enterprise_service import EnterpriseService from services.entities.agent_entities import ComposerSavePayload, RosterListQuery from services.feature_service import FeatureService +from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task AgentPublicationStatus = Literal["published", "drafts"] @@ -684,6 +687,15 @@ class AgentAppListApi(Resource): ) app = AppService().create_app(current_tenant_id, params, current_user, session=session) + if dify_config.RBAC_ENABLED: + enterprise_rbac_service.RBACService.AppAccess.replace_whitelist( + current_tenant_id, + current_user.id, + str(app.id), + enterprise_rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=True), + ) + initialize_created_app_rbac_access_task.delay(current_tenant_id, current_user.id, app_id=app.id) + return _serialize_agent_app_detail(session, app, current_user=current_user), 201 diff --git a/api/services/enterprise/rbac_service.py b/api/services/enterprise/rbac_service.py index b9a6db00e00..b3d1bfac0fd 100644 --- a/api/services/enterprise/rbac_service.py +++ b/api/services/enterprise/rbac_service.py @@ -475,6 +475,7 @@ _LEGACY_WORKSPACE_DATASET_OPERATOR_KEYS: list[str] = [ _LEGACY_APP_OWNER_KEYS: list[str] = [ "app.acl.preview", + "app.acl.access_point_manage", "app.acl.view_layout", "app.acl.test_and_run", "app.acl.edit", @@ -490,6 +491,7 @@ _LEGACY_APP_OWNER_KEYS: list[str] = [ _LEGACY_APP_ADMIN_KEYS: list[str] = [ "app.acl.preview", "app.acl.view_layout", + "app.acl.access_point_manage", "app.acl.test_and_run", "app.acl.edit", "app.acl.import_export_dsl", @@ -504,6 +506,7 @@ _LEGACY_APP_ADMIN_KEYS: list[str] = [ _LEGACY_APP_EDITOR_KEYS: list[str] = [ "app.acl.preview", + "app.acl.access_point_manage", "app.acl.view_layout", "app.acl.test_and_run", "app.acl.edit", diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py index c2b43d8cb22..27a543fd2a2 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from datetime import datetime from inspect import getsource, unwrap from types import SimpleNamespace @@ -308,9 +309,15 @@ def account_id() -> str: def test_agent_app_list_and_create_use_agent_route( - app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str, sqlite_session: Session + app: Flask, + monkeypatch: pytest.MonkeyPatch, + account_id: str, + sqlite_session: Session, + config_overrides: Callable[..., None], ) -> None: captured: dict[str, object] = {} + replace_whitelist = MagicMock() + initialize_access = MagicMock() class FakeAppService: def get_app(self, app_obj: object, *, session: object) -> object: @@ -396,7 +403,9 @@ def test_agent_app_list_and_create_use_agent_route( lambda _self, **kwargs: {"agent-list": "debug-conversation-list"}, ) monkeypatch.setattr( - roster_controller.AgentRosterService, "count_agent_app_debug_conversation_messages", lambda _self, **kwargs: 0 + roster_controller.AgentRosterService, + "count_agent_app_debug_conversation_messages", + lambda _self, **kwargs: 0, ) def get_or_create_debug_conversation(_self: object, **kwargs: object) -> str: @@ -413,6 +422,13 @@ def test_agent_app_list_and_create_use_agent_route( "get_system_features", lambda: SimpleNamespace(webapp_auth=SimpleNamespace(enabled=False)), ) + config_overrides(RBAC_ENABLED=True) + monkeypatch.setattr( + roster_controller.enterprise_rbac_service.RBACService.AppAccess, + "replace_whitelist", + replace_whitelist, + ) + monkeypatch.setattr(roster_controller.initialize_created_app_rbac_access_task, "delay", initialize_access) with app.test_request_context( "/console/api/agent?page=1&limit=10&mode=workflow&sort_by=recently_created" "&is_created_by_me=true&publication_status=published" @@ -453,12 +469,22 @@ def test_agent_app_list_and_create_use_agent_route( assert count_params.agent_is_published is True with app.test_request_context( "/console/api/agent", - json={"name": "Iris", "description": "Agent app", "role": "Coordinator", "icon_type": "emoji", "icon": "robot"}, + json={ + "name": "Iris", + "description": "Agent app", + "role": "Coordinator", + "icon_type": "emoji", + "icon": "robot", + }, ): created, status = unwrap(AgentAppListApi.post)( AgentAppListApi(), AgentAppCreatePayload( - name="Iris", description="Agent app", role="Coordinator", icon_type="emoji", icon="robot" + name="Iris", + description="Agent app", + role="Coordinator", + icon_type="emoji", + icon="robot", ), sqlite_session, "tenant-1", @@ -481,6 +507,81 @@ def test_agent_app_list_and_create_use_agent_route( "account_id": account_id, "commit": False, } + replace_whitelist.assert_called_once() + assert replace_whitelist.call_args.args[:3] == ("tenant-1", account_id, "app-created") + replace_payload = replace_whitelist.call_args.args[3] + assert replace_payload.automatic_include_workspace_members is True + initialize_access.assert_called_once_with("tenant-1", account_id, app_id="app-created") + + +def test_agent_app_create_skips_rbac_access_initialization_when_rbac_is_disabled( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + account_id: str, + sqlite_session: Session, + config_overrides: Callable[..., None], +) -> None: + replace_whitelist = MagicMock() + initialize_access = MagicMock() + + class FakeAppService: + def get_app(self, app_obj: object, *, session: object) -> object: + return app_obj + + def create_app(self, tenant_id: str, params, current_user: object, *, session: object) -> object: + return _app_detail_obj(id="app-created", bound_agent_id="agent-created") + + monkeypatch.setattr(roster_controller, "AppService", FakeAppService) + monkeypatch.setattr( + roster_controller.AgentRosterService, + "get_app_backing_agent", + lambda _self, **kwargs: Agent( + id="agent-created", + app_id="app-created", + backing_app_id=None, + role="Created role", + active_config_snapshot_id=None, + ), + ) + monkeypatch.setattr( + roster_controller.AgentRosterService, + "get_or_create_build_conversation", + lambda _self, **kwargs: "debug-conversation-created", + ) + monkeypatch.setattr( + roster_controller.AgentRosterService, "count_agent_app_debug_conversation_messages", lambda _self, **kwargs: 0 + ) + monkeypatch.setattr( + roster_controller.FeatureService, + "get_system_features", + lambda: SimpleNamespace(webapp_auth=SimpleNamespace(enabled=False)), + ) + config_overrides(RBAC_ENABLED=False) + monkeypatch.setattr( + roster_controller.enterprise_rbac_service.RBACService.AppAccess, + "replace_whitelist", + replace_whitelist, + ) + monkeypatch.setattr(roster_controller.initialize_created_app_rbac_access_task, "delay", initialize_access) + + with app.test_request_context( + "/console/api/agent", + json={"name": "Iris", "description": "Agent app", "role": "Coordinator", "icon_type": "emoji", "icon": "robot"}, + ): + created, status = unwrap(AgentAppListApi.post)( + AgentAppListApi(), + AgentAppCreatePayload( + name="Iris", description="Agent app", role="Coordinator", icon_type="emoji", icon="robot" + ), + sqlite_session, + "tenant-1", + _account(account_id=account_id), + ) + + assert status == 201 + assert created["id"] == "agent-created" + replace_whitelist.assert_not_called() + initialize_access.assert_not_called() def test_agent_app_create_payload_allows_optional_role() -> None: diff --git a/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/__tests__/layout-main.spec.tsx b/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/__tests__/layout-main.spec.tsx index 8c6acc78866..9a52e32d0e7 100644 --- a/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/__tests__/layout-main.spec.tsx +++ b/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/__tests__/layout-main.spec.tsx @@ -238,9 +238,11 @@ describe('AppDetailLayout', () => { expect(useStore.getState().appDetail?.id).toBe('app-1') }) - it('should allow access point pages without app deploy or app ACL permissions', async () => { + it('should allow users with access point permission to open access point directly', async () => { mockPathname = '/app/app-1/access-point' - mockFetchAppDetailDirect.mockResolvedValue(createAppDetail({ permission_keys: [] })) + mockFetchAppDetailDirect.mockResolvedValue( + createAppDetail({ permission_keys: [AppACLPermission.AccessPoint] }), + ) render( @@ -254,6 +256,44 @@ describe('AppDetailLayout', () => { expect(useStore.getState().appDetail?.id).toBe('app-1') }) + it('should redirect access point pages when access point permission is missing', async () => { + mockPathname = '/app/app-1/access-point' + mockFetchAppDetailDirect.mockResolvedValue( + createAppDetail({ permission_keys: [AppACLPermission.Monitor] }), + ) + + render( + +
App page content
+ , + ) + + await waitFor(() => { + expect(mockReplace).toHaveBeenCalledWith('/app/app-1/overview') + }) + expect(screen.queryByText('App page content')).not.toBeInTheDocument() + expect(useStore.getState().appDetail).toBeUndefined() + }) + + it('should keep access point content hidden while redirecting cached app data without permission', async () => { + mockPathname = '/app/app-1/access-point' + useStore + .getState() + .setAppDetail(createAppDetail({ permission_keys: [AppACLPermission.Monitor] })) + + render( + +
App page content
+ , + ) + + expect(screen.queryByText('App page content')).not.toBeInTheDocument() + await waitFor(() => { + expect(mockReplace).toHaveBeenCalledWith('/app/app-1/overview') + }) + expect(mockFetchAppDetailDirect).not.toHaveBeenCalled() + }) + it('should redirect deploy pages when app deploy ACL permission is missing', async () => { mockPathname = '/app/app-1/deploy' mockFetchAppDetailDirect.mockResolvedValue( @@ -317,7 +357,7 @@ describe('AppDetailLayout', () => { ) await waitFor(() => { - expect(mockReplace).toHaveBeenCalledWith('/app/app-1/access-point') + expect(mockReplace).toHaveBeenCalledWith('/apps') }) expect(screen.queryByText('App page content')).not.toBeInTheDocument() expect(useStore.getState().appDetail).toBeUndefined() @@ -488,7 +528,7 @@ describe('AppDetailLayout', () => { ) await waitFor(() => { - expect(mockReplace).toHaveBeenCalledWith('/app/app-1/access-point') + expect(mockReplace).toHaveBeenCalledWith('/apps') }) expect(screen.queryByText('App page content')).not.toBeInTheDocument() expect(useStore.getState().appDetail).toBeUndefined() diff --git a/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/layout-main.tsx b/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/layout-main.tsx index 7f9f6973ca8..df55021f289 100644 --- a/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/layout-main.tsx +++ b/web/app/(commonLayout)/app/(appDetailLayout)/[appId]/layout-main.tsx @@ -78,8 +78,28 @@ const AppDetailLayout: FC = (props) => { appDetail?.id === appId ? appDetail : appDetailRes?.id === appId ? appDetailRes : null const pageTitle = appDetailPageTitle(pathname, t) const appName = routeAppDetail?.id === appId ? routeAppDetail.name : undefined + const isAppACLContextReady = + !!routeAppDetail && + !!currentWorkspace.id && + !isLoadingCurrentWorkspace && + !isLoadingWorkspacePermissionKeys && + !isLoadingAppDetail + const appACLCapabilities = React.useMemo( + () => + routeAppDetail && isAppACLContextReady + ? getAppACLCapabilities(routeAppDetail.permission_keys, { + currentUserId, + resourceMaintainer: routeAppDetail.maintainer, + workspacePermissionKeys, + isRbacEnabled, + }) + : null, + [currentUserId, isAppACLContextReady, isRbacEnabled, routeAppDetail, workspacePermissionKeys], + ) const shouldBlockAgentResourceAccess = routeAppDetail?.mode === AppModeEnum.AGENT && pathname.endsWith('/access-config') + const shouldBlockAccessPointAccess = + pathname.endsWith('/access-point') && !appACLCapabilities?.canAccessPoint useDocumentTitle(`${pageTitle} · ${appName || t(($) => $['menus.appDetail'], { ns: 'common' })}`) @@ -120,28 +140,16 @@ const AppDetailLayout: FC = (props) => { }, [appId, router, setAppDetail]) useEffect(() => { - if ( - !routeAppDetail || - !currentWorkspace.id || - isLoadingCurrentWorkspace || - isLoadingWorkspacePermissionKeys || - isLoadingAppDetail - ) - return + if (!routeAppDetail || !isAppACLContextReady || !appACLCapabilities) return if (routeAppDetail.id !== appId) return - const appACLCapabilities = getAppACLCapabilities(routeAppDetail.permission_keys, { - currentUserId, - resourceMaintainer: routeAppDetail.maintainer, - workspacePermissionKeys, - isRbacEnabled, - }) const isLayoutPath = pathname.endsWith('configuration') || pathname.endsWith('workflow') const isLogsPath = pathname.endsWith('logs') const isAnnotationsPath = pathname.endsWith('annotations') const isOverviewPath = pathname.endsWith('overview') const isAccessConfigPath = pathname.endsWith('access-config') const isDeployPath = pathname.endsWith('deploy') + const isAccessPointPath = pathname.endsWith('access-point') if ( (isLayoutPath && !appACLCapabilities.canAccessLayout) || (isLogsPath && !appACLCapabilities.canAccessLogAndAnnotation) || @@ -150,7 +158,8 @@ const AppDetailLayout: FC = (props) => { (isAccessConfigPath && (routeAppDetail.mode === AppModeEnum.AGENT || !appACLCapabilities.canAccessConfig)) || (isDeployPath && - (routeAppDetail.mode !== AppModeEnum.WORKFLOW || !appACLCapabilities.canDeploy)) + (routeAppDetail.mode !== AppModeEnum.WORKFLOW || !appACLCapabilities.canDeploy)) || + (isAccessPointPath && !appACLCapabilities.canAccessPoint) ) { router.replace( getRedirectionPath(routeAppDetail, { @@ -180,14 +189,12 @@ const AppDetailLayout: FC = (props) => { if (appDetailRes && appDetail?.id !== appDetailRes.id) setAppDetail({ ...appDetailRes, enable_sso: false }) }, [ + appACLCapabilities, appDetail?.id, appDetailRes, appId, currentUserId, - currentWorkspace.id, - isLoadingAppDetail, - isLoadingCurrentWorkspace, - isLoadingWorkspacePermissionKeys, + isAppACLContextReady, isRbacEnabled, pathname, routeAppDetail, @@ -198,7 +205,7 @@ const AppDetailLayout: FC = (props) => { const isWorkflowPage = pathname.endsWith('/workflow') const content = - !appDetail || shouldBlockAgentResourceAccess ? ( + !appDetail || shouldBlockAgentResourceAccess || shouldBlockAccessPointAccess ? (
diff --git a/web/app/components/app-sidebar/__tests__/app-detail-section.spec.tsx b/web/app/components/app-sidebar/__tests__/app-detail-section.spec.tsx index e84d925dbdc..0ff08042648 100644 --- a/web/app/components/app-sidebar/__tests__/app-detail-section.spec.tsx +++ b/web/app/components/app-sidebar/__tests__/app-detail-section.spec.tsx @@ -186,7 +186,10 @@ describe('AppDetailSection', () => { ).not.toBeInTheDocument() }) - it('should render access point navigation using its app route', () => { + it('should render access point navigation when access point permission is granted', () => { + // Arrange + mockAppPermissionKeys = [AppACLPermission.AccessPoint] + // Act render() @@ -200,6 +203,19 @@ describe('AppDetailSection', () => { ).not.toBeInTheDocument() }) + it('should hide access point navigation when access point permission is missing', () => { + // Arrange + mockAppPermissionKeys = [AppACLPermission.Monitor] + + // Act + render() + + // Assert + expect( + screen.queryByRole('link', { name: 'common.appMenus.accessPoint' }), + ).not.toBeInTheDocument() + }) + it('should render deploy navigation with app deploy ACL regardless of the legacy workspace role', () => { // Arrange mockAppMode = 'workflow' diff --git a/web/app/components/app-sidebar/app-detail-section.tsx b/web/app/components/app-sidebar/app-detail-section.tsx index 3176e2cac21..677ac50aea4 100644 --- a/web/app/components/app-sidebar/app-detail-section.tsx +++ b/web/app/components/app-sidebar/app-detail-section.tsx @@ -120,12 +120,16 @@ const AppDetailSection = ({ expand = true }: AppDetailSectionProps) => { }, ] : []), - { - name: t(($) => $['appMenus.accessPoint'], { ns: 'common' }), - href: `/app/${appId}/access-point`, - icon: accessPointNavIcon, - selectedIcon: accessPointNavIcon, - }, + ...(appACLCapabilities.canAccessPoint + ? [ + { + name: t(($) => $['appMenus.accessPoint'], { ns: 'common' }), + href: `/app/${appId}/access-point`, + icon: accessPointNavIcon, + selectedIcon: accessPointNavIcon, + }, + ] + : []), ...(supportsAppDeploy && appACLCapabilities.canDeploy ? [ { diff --git a/web/app/components/app/app-publisher/__tests__/environment-deployment-flow.spec.tsx b/web/app/components/app/app-publisher/__tests__/environment-deployment-flow.spec.tsx index 8aa12ed5457..9c80ad25f11 100644 --- a/web/app/components/app/app-publisher/__tests__/environment-deployment-flow.spec.tsx +++ b/web/app/components/app/app-publisher/__tests__/environment-deployment-flow.spec.tsx @@ -295,6 +295,7 @@ function renderFlow( return render( ({ default: ({ @@ -314,6 +315,7 @@ describe('app-publisher sections', () => { description: 'Workflow description', }} appURL="https://example.com/app" + canAccessPoint disabledFunctionButton={false} disabledFunctionTooltip="disabled" handleOpenRunConfig={handleOpenRunConfig} @@ -494,6 +496,7 @@ describe('app-publisher sections', () => { mode: AppModeEnum.WORKFLOW, }} appURL="https://example.com/app" + canAccessPoint disabledFunctionButton={false} hasHumanInputNode={false} hasTriggerNode @@ -517,11 +520,49 @@ describe('app-publisher sections', () => { ) }) + it('should hide the built-in Access Point action without permission', () => { + render( + , + ) + + expect(screen.queryByText(/(?:^|\.)appMenus\.accessPoint(?=$|:)/)).not.toBeInTheDocument() + expect(screen.getByRole('link', { name: /appMenus\.deploy\b/ })).toHaveAttribute( + 'href', + '/app/workflow-app/deploy', + ) + }) + + it('should hide the environment Access Point action without permission', () => { + render( + , + ) + + expect(screen.queryByText(/(?:^|\.)appMenus\.accessPoint(?=$|:)/)).not.toBeInTheDocument() + expect(screen.getByText(/(?:^|\.)appMenus\.deploy(?=$|:)/)).toBeInTheDocument() + }) + it('should expose unavailable quick links as disabled buttons before the first publish', () => { render( void @@ -41,6 +42,7 @@ type PublisherActionsSectionProps = Pick< export function PublisherActionsSection({ appDetail, appURL, + canAccessPoint = false, disabledFunctionButton, disabledFunctionTooltip, handleOpenRunConfig, @@ -114,14 +116,16 @@ export function PublisherActionsSection({ {disabledFunctionTooltip} )} - $['common.accessPointDescription'], { ns: 'workflow' })} - link={appId ? `/app/${appId}/access-point` : undefined} - icon={} - > - {t(($) => $['appMenus.accessPoint'], { ns: 'common' })} - + {canAccessPoint && ( + $['common.accessPointDescription'], { ns: 'workflow' })} + link={appId ? `/app/${appId}/access-point` : undefined} + icon={} + > + {t(($) => $['appMenus.accessPoint'], { ns: 'common' })} + + )} {showDeploy && ( - $['common.accessPointDescription'], { ns: 'workflow' })} - link={accessPointHref} - icon={} - > - {t(($) => $['appMenus.accessPoint'], { ns: 'common' })} - + {canAccessPoint && ( + $['common.accessPointDescription'], { ns: 'workflow' })} + link={accessPointHref} + icon={} + > + {t(($) => $['appMenus.accessPoint'], { ns: 'common' })} + + )} $['common.deployDescription'], { ns: 'workflow' })} diff --git a/web/app/components/app/app-publisher/environment-deployment-flow/index.tsx b/web/app/components/app/app-publisher/environment-deployment-flow/index.tsx index dce08832bc0..9f0b1201993 100644 --- a/web/app/components/app/app-publisher/environment-deployment-flow/index.tsx +++ b/web/app/components/app/app-publisher/environment-deployment-flow/index.tsx @@ -15,6 +15,7 @@ import { PublisherEnvironmentSummarySection } from './summary-section' type PublisherEnvironmentFlowProps = { appId?: string + canAccessPoint?: boolean deployment?: EnvironmentDeployment environmentId: string environmentName: string @@ -28,6 +29,7 @@ type PublisherEnvironmentFlowProps = { export function PublisherEnvironmentFlow({ appId, + canAccessPoint = false, deployment, environmentId, environmentName, @@ -91,6 +93,7 @@ export function PublisherEnvironmentFlow({ /> diff --git a/web/app/components/app/app-publisher/index.tsx b/web/app/components/app/app-publisher/index.tsx index 48278b77cfc..38e9b724d1e 100644 --- a/web/app/components/app/app-publisher/index.tsx +++ b/web/app/components/app/app-publisher/index.tsx @@ -17,12 +17,13 @@ export function AppPublisher(props: AppPublisherProps) { select: (data) => data.profile.id, }) const workspacePermissionKeys = useAtomValue(workspacePermissionKeysAtom) - const canDeploy = getAppACLCapabilities(appDetail?.permission_keys, { + const appACLCapabilities = getAppACLCapabilities(appDetail?.permission_keys, { currentUserId, resourceMaintainer: appDetail?.maintainer, workspacePermissionKeys, - }).canDeploy - const supportsMultiEnvironment = appDetail?.mode === AppModeEnum.WORKFLOW && canDeploy + }) + const supportsMultiEnvironment = + appDetail?.mode === AppModeEnum.WORKFLOW && appACLCapabilities.canDeploy return ( void } export function PublisherContent({ + canAccessPoint, crossAxisOffset = 0, debugWithMultipleModel = false, disabled = false, @@ -212,6 +214,7 @@ export function PublisherContent({ actions: { appDetail, appURL, + canAccessPoint, disabledFunctionButton, disabledFunctionTooltip, handleOpenRunConfig: workflowLaunch.openDialog, @@ -236,6 +239,7 @@ export function PublisherContent({ disabled={disabled} environmentPublisher={{ appId: appDetail?.id, + canAccessPoint, deployment: selectedEnvironmentDeployment, environmentId: selectedEnvironmentId, environmentName: diff --git a/web/app/components/app/deploy/__tests__/index.spec.tsx b/web/app/components/app/deploy/__tests__/index.spec.tsx index 11888a09817..79c1f082294 100644 --- a/web/app/components/app/deploy/__tests__/index.spec.tsx +++ b/web/app/components/app/deploy/__tests__/index.spec.tsx @@ -650,7 +650,7 @@ function render( return renderWithConsoleQuery(ui, { queryClient }) } -let appPermissionKeys: string[] = [AppACLPermission.Deploy] +let appPermissionKeys: string[] = [AppACLPermission.AccessPoint, AppACLPermission.Deploy] let appDetailAvailable = true const mockConsoleState = vi.hoisted(() => ({ workspacePermissionKeys: [] as string[], @@ -744,7 +744,7 @@ vi.mock('@langgenius/dify-ui/toast', () => ({ describe('AppDeploy', () => { beforeEach(() => { vi.clearAllMocks() - appPermissionKeys = [AppACLPermission.Deploy] + appPermissionKeys = [AppACLPermission.AccessPoint, AppACLPermission.Deploy] appDetailAvailable = true mockBuiltInEnvironment.appDetail.enable_api = false mockBuiltInEnvironment.appDetail.enable_site = true @@ -815,6 +815,18 @@ describe('AppDeploy', () => { ).toHaveAttribute('href', '/app/app-1/access-point?environment=canary&accessPoint=serviceApi') }) + it('keeps active access points non-navigable without access point permission', () => { + appPermissionKeys = [AppACLPermission.Deploy] + + render() + + const canaryRow = within(screen.getByRole('row', { name: /Canary/ })) + const webAppLabel = + 'agentV2.agentDetail.access.webApp.title · agentV2.agentDetail.access.status.inService' + expect(canaryRow.queryByRole('link', { name: webAppLabel })).not.toBeInTheDocument() + expect(canaryRow.getByRole('button', { name: webAppLabel })).toBeDisabled() + }) + it('renders the built-in version, access points, and publisher from live app data', () => { render() diff --git a/web/app/components/app/deploy/built-in-environment-card/index.tsx b/web/app/components/app/deploy/built-in-environment-card/index.tsx index d04dce259a8..607250e7ce6 100644 --- a/web/app/components/app/deploy/built-in-environment-card/index.tsx +++ b/web/app/components/app/deploy/built-in-environment-card/index.tsx @@ -19,7 +19,7 @@ function Divider() { return
} -export function BuiltInEnvironmentCard() { +export function BuiltInEnvironmentCard({ canAccessPoint = false }: { canAccessPoint?: boolean }) { const { t } = useTranslation('deployments') const { formatTime } = useTimestamp() const appDetail = useAppStore((state) => state.appDetail) @@ -90,7 +90,9 @@ export function BuiltInEnvironmentCard() { key={accessPoint} accessPoint={accessPoint} active={activeAccessPoints[accessPoint]} - href={getAccessPointHref(appId, 'built-in', accessPoint)} + href={ + canAccessPoint ? getAccessPointHref(appId, 'built-in', accessPoint) : undefined + } /> ))}
diff --git a/web/app/components/app/deploy/environment-table/index.tsx b/web/app/components/app/deploy/environment-table/index.tsx index cbee27225ed..cc9b9451e59 100644 --- a/web/app/components/app/deploy/environment-table/index.tsx +++ b/web/app/components/app/deploy/environment-table/index.tsx @@ -23,6 +23,7 @@ import { EnvironmentRow } from './row' type EnvironmentTableProps = { appId: string + canAccessPoint?: boolean onChangeVersion?: (deployment: EnvironmentDeployment) => void onDeployLatest?: (deployment: EnvironmentDeployment) => void onDeployToEnvironment?: (environment: AppEnvironment) => void @@ -32,6 +33,7 @@ type EnvironmentTableProps = { export function EnvironmentTable({ appId, + canAccessPoint = false, onChangeVersion, onDeployLatest, onDeployToEnvironment, @@ -132,6 +134,7 @@ export function EnvironmentTable({ void onDeployLatest?: (deployment: EnvironmentDeployment) => void @@ -60,7 +62,11 @@ export function EnvironmentRow({ key={accessPoint} accessPoint={accessPoint} active={isAccessPointActive(accessPoint)} - href={getAccessPointHref(appId, row.environment.id, accessPoint)} + href={ + canAccessPoint + ? getAccessPointHref(appId, row.environment.id, accessPoint) + : undefined + } /> ))} diff --git a/web/app/components/app/deploy/index.tsx b/web/app/components/app/deploy/index.tsx index c22034584dd..4f1bee0ed1e 100644 --- a/web/app/components/app/deploy/index.tsx +++ b/web/app/components/app/deploy/index.tsx @@ -22,7 +22,7 @@ import { useRefreshAppEnvironmentsAfterDeploymentPolling } from './use-refresh-a import { useUndeployWorkflow } from './use-undeploy-workflow' import { toDeploymentVersion } from './version' -function AppDeployContent({ appId }: { appId: string }) { +function AppDeployContent({ appId, canAccessPoint }: { appId: string; canAccessPoint: boolean }) { const { t } = useTranslation('deployments') const { t: tCommon } = useTranslation('common') const { t: tWorkflow } = useTranslation('workflow') @@ -86,9 +86,10 @@ function AppDeployContent({ appId }: { appId: string }) {
- + setDeploymentRequest({ environment: environment.display_name, @@ -139,17 +140,17 @@ export default function AppDeploy() { if (!appDetail) return - const canDeploy = getAppACLCapabilities(appDetail.permission_keys, { + const appACLCapabilities = getAppACLCapabilities(appDetail.permission_keys, { currentUserId, resourceMaintainer: appDetail.maintainer, workspacePermissionKeys, - }).canDeploy + }) - if (appDetail.mode !== AppModeEnum.WORKFLOW || !canDeploy) return null + if (appDetail.mode !== AppModeEnum.WORKFLOW || !appACLCapabilities.canDeploy) return null return ( - + ) } diff --git a/web/app/components/app/deploy/shared/__tests__/access-point-icon.spec.tsx b/web/app/components/app/deploy/shared/__tests__/access-point-icon.spec.tsx new file mode 100644 index 00000000000..80d18eaaee6 --- /dev/null +++ b/web/app/components/app/deploy/shared/__tests__/access-point-icon.spec.tsx @@ -0,0 +1,35 @@ +import { screen } from '@testing-library/react' +import { renderWithConsoleQuery as render } from '@/test/console/query-data' +import { AccessPointIcon } from '../access-point-icon' + +describe('AccessPointIcon', () => { + it('links active access points when navigation is allowed', () => { + render( + , + ) + + expect(screen.getByRole('link')).toHaveAttribute( + 'href', + '/app/app-1/access-point?environment=built-in&accessPoint=webApp', + ) + }) + + it('keeps active access points visually active when navigation is not allowed', () => { + render() + + expect(screen.queryByRole('link')).not.toBeInTheDocument() + expect(screen.getByRole('button')).toBeDisabled() + expect(screen.getByRole('button')).not.toHaveClass('opacity-30') + }) + + it('dims inactive access points', () => { + render() + + expect(screen.getByRole('button')).toBeDisabled() + expect(screen.getByRole('button')).toHaveClass('opacity-30') + }) +}) diff --git a/web/app/components/app/deploy/shared/access-point-icon.tsx b/web/app/components/app/deploy/shared/access-point-icon.tsx index 734e9b694ee..253429f03fc 100644 --- a/web/app/components/app/deploy/shared/access-point-icon.tsx +++ b/web/app/components/app/deploy/shared/access-point-icon.tsx @@ -32,7 +32,7 @@ export function AccessPointIcon({ }: { active: boolean accessPoint: AccessPoint - href: string + href?: string }) { const { t } = useTranslation('agentV2') const labels = useAccessPointLabels() @@ -40,9 +40,11 @@ export function AccessPointIcon({ ? t(($) => $['agentDetail.access.status.inService']) : t(($) => $['agentDetail.access.status.outOfService']) const label = `${labels[accessPoint]} · ${status}` + const canNavigate = active && Boolean(href) const triggerClassName = cn( 'flex size-5 shrink-0 items-center justify-center rounded-md border border-divider-regular text-text-secondary shadow-xs outline-hidden focus-visible:ring-2 focus-visible:ring-state-accent-solid', - active ? 'cursor-pointer hover:bg-state-base-hover' : 'cursor-not-allowed opacity-30', + active && (canNavigate ? 'cursor-pointer hover:bg-state-base-hover' : 'cursor-default'), + !active && 'cursor-not-allowed opacity-30', ) const icon = ( @@ -52,7 +54,7 @@ export function AccessPointIcon({ {icon} diff --git a/web/app/components/header/account-setting/access-rules-page/permission-set-modal/__tests__/index.spec.tsx b/web/app/components/header/account-setting/access-rules-page/permission-set-modal/__tests__/index.spec.tsx index 7c4b00da20e..764b51cb52c 100644 --- a/web/app/components/header/account-setting/access-rules-page/permission-set-modal/__tests__/index.spec.tsx +++ b/web/app/components/header/account-setting/access-rules-page/permission-set-modal/__tests__/index.spec.tsx @@ -19,6 +19,7 @@ const expectedAppACLPermissionKeys = [ 'app.acl.tracing_config', 'app.acl.log_and_annotation', 'app.acl.access_config', + 'app.acl.access_point_manage', ] const getPermissionKeyMatcher = (permissionKey: string) => diff --git a/web/features/home/continue-work/__tests__/item.spec.tsx b/web/features/home/continue-work/__tests__/item.spec.tsx index ee61d10d6ce..37b74a4b944 100644 --- a/web/features/home/continue-work/__tests__/item.spec.tsx +++ b/web/features/home/continue-work/__tests__/item.spec.tsx @@ -168,10 +168,15 @@ describe('ContinueWorkItem', () => { ) }) - it('should fall back to access point when RBAC is disabled for an access-config-only app', () => { - renderItem(createApp({ permission_keys: [AppACLPermission.AccessConfig] }), { - rbac_enabled: false, - }) + it('should fall back to access point when RBAC is disabled for an access-config app with access point permission', () => { + renderItem( + createApp({ + permission_keys: [AppACLPermission.AccessConfig, AppACLPermission.AccessPoint], + }), + { + rbac_enabled: false, + }, + ) expect(screen.getByRole('link', { name: /Continue App/ })).toHaveAttribute( 'href', diff --git a/web/i18n/ar-TN/permission-keys.json b/web/i18n/ar-TN/permission-keys.json index a7f129455c5..18321e3b376 100644 --- a/web/i18n/ar-TN/permission-keys.json +++ b/web/i18n/ar-TN/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "إدارة إعدادات امتداد API", "app.access_config": "تكوين أذونات الوصول إلى التطبيق", "app.acl.access_config": "عرض أذونات الوصول وإدارتها", + "app.acl.access_point_manage": "عرض نقاط الوصول وإدارتها", "app.acl.delete": "حذف التطبيق", "app.acl.deploy": "نشر التطبيق", "app.acl.edit": "تعديل معلومات التطبيق وتنسيقه", diff --git a/web/i18n/de-DE/permission-keys.json b/web/i18n/de-DE/permission-keys.json index fb32d92c699..2d546c0e082 100644 --- a/web/i18n/de-DE/permission-keys.json +++ b/web/i18n/de-DE/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "API-Erweiterungskonfiguration verwalten", "app.access_config": "App-Zugriffsberechtigungen konfigurieren", "app.acl.access_config": "Zugriffsberechtigungen anzeigen und verwalten", + "app.acl.access_point_manage": "Zugangspunkte anzeigen und verwalten", "app.acl.delete": "App löschen", "app.acl.deploy": "App bereitstellen", "app.acl.edit": "App-Informationen bearbeiten und App orchestrieren", diff --git a/web/i18n/en-US/permission-keys.json b/web/i18n/en-US/permission-keys.json index 2344caa11e1..e69f68bb0b3 100644 --- a/web/i18n/en-US/permission-keys.json +++ b/web/i18n/en-US/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Manage API extension configuration", "app.access_config": "Configure app access permissions", "app.acl.access_config": "View and manage access permissions", + "app.acl.access_point_manage": "View and manage access points", "app.acl.delete": "Delete app", "app.acl.deploy": "Deploy app", "app.acl.edit": "Edit app information and orchestrate app", diff --git a/web/i18n/es-ES/permission-keys.json b/web/i18n/es-ES/permission-keys.json index af3336ceea7..d7d5e9d57c8 100644 --- a/web/i18n/es-ES/permission-keys.json +++ b/web/i18n/es-ES/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Gestionar la configuración de la extensión de API", "app.access_config": "Configurar los permisos de acceso de la app", "app.acl.access_config": "Ver y gestionar los permisos de acceso", + "app.acl.access_point_manage": "Ver y gestionar los puntos de acceso", "app.acl.delete": "Eliminar app", "app.acl.deploy": "Desplegar la app", "app.acl.edit": "Editar la información y orquestar la app", diff --git a/web/i18n/fa-IR/permission-keys.json b/web/i18n/fa-IR/permission-keys.json index 372374ed3b8..5e739bed336 100644 --- a/web/i18n/fa-IR/permission-keys.json +++ b/web/i18n/fa-IR/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "مدیریت پیکربندی افزونه API", "app.access_config": "پیکربندی مجوزهای دسترسی برنامه", "app.acl.access_config": "مشاهده و مدیریت مجوزهای دسترسی", + "app.acl.access_point_manage": "مشاهده و مدیریت نقاط دسترسی", "app.acl.delete": "حذف برنامه", "app.acl.deploy": "استقرار برنامه", "app.acl.edit": "ویرایش اطلاعات برنامه و هماهنگ‌سازی برنامه", diff --git a/web/i18n/fr-FR/permission-keys.json b/web/i18n/fr-FR/permission-keys.json index 074da133fae..c547052586d 100644 --- a/web/i18n/fr-FR/permission-keys.json +++ b/web/i18n/fr-FR/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Gérer la configuration de l'extension API", "app.access_config": "Configurer les autorisations d'accès à l'application", "app.acl.access_config": "Afficher et gérer les autorisations d'accès", + "app.acl.access_point_manage": "Afficher et gérer les points d’accès", "app.acl.delete": "Supprimer l'application", "app.acl.deploy": "Déployer l'application", "app.acl.edit": "Modifier les informations et orchestrer l'application", diff --git a/web/i18n/hi-IN/permission-keys.json b/web/i18n/hi-IN/permission-keys.json index 314cbfba84b..0779d870342 100644 --- a/web/i18n/hi-IN/permission-keys.json +++ b/web/i18n/hi-IN/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "API एक्सटेंशन कॉन्फ़िगरेशन प्रबंधित करें", "app.access_config": "ऐप एक्सेस अनुमतियाँ कॉन्फ़िगर करें", "app.acl.access_config": "एक्सेस अनुमतियाँ देखें और प्रबंधित करें", + "app.acl.access_point_manage": "एक्सेस पॉइंट देखें और प्रबंधित करें", "app.acl.delete": "ऐप हटाएं", "app.acl.deploy": "ऐप डिप्लॉय करें", "app.acl.edit": "ऐप की जानकारी संपादित करें और ऐप को ऑर्केस्ट्रेट करें", diff --git a/web/i18n/id-ID/permission-keys.json b/web/i18n/id-ID/permission-keys.json index 4b04ea96f00..349bff75538 100644 --- a/web/i18n/id-ID/permission-keys.json +++ b/web/i18n/id-ID/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Kelola konfigurasi ekstensi API", "app.access_config": "Konfigurasikan izin akses aplikasi", "app.acl.access_config": "Lihat dan kelola izin akses", + "app.acl.access_point_manage": "Lihat dan kelola titik akses", "app.acl.delete": "Hapus aplikasi", "app.acl.deploy": "Deploy aplikasi", "app.acl.edit": "Edit informasi aplikasi dan orkestrasikan aplikasi", diff --git a/web/i18n/it-IT/permission-keys.json b/web/i18n/it-IT/permission-keys.json index 899adce084b..b5f3ebf8094 100644 --- a/web/i18n/it-IT/permission-keys.json +++ b/web/i18n/it-IT/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Gestisci la configurazione delle estensioni API", "app.access_config": "Configura i permessi di accesso all'app", "app.acl.access_config": "Visualizza e gestisci i permessi di accesso", + "app.acl.access_point_manage": "Visualizza e gestisci i punti di accesso", "app.acl.delete": "Elimina app", "app.acl.deploy": "Distribuisci app", "app.acl.edit": "Modifica le informazioni e orchestra l'app", diff --git a/web/i18n/ja-JP/permission-keys.json b/web/i18n/ja-JP/permission-keys.json index 1b0c567f0e8..53033e2dc10 100644 --- a/web/i18n/ja-JP/permission-keys.json +++ b/web/i18n/ja-JP/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "API拡張設定を管理", "app.access_config": "アプリアクセス権限を設定", "app.acl.access_config": "アクセス権限の表示と管理", + "app.acl.access_point_manage": "アクセスポイントの表示と管理", "app.acl.delete": "アプリを削除", "app.acl.deploy": "アプリをデプロイ", "app.acl.edit": "アプリ情報の編集とアプリのオーケストレーション", diff --git a/web/i18n/ko-KR/permission-keys.json b/web/i18n/ko-KR/permission-keys.json index 3a9981a602a..45517f6acb9 100644 --- a/web/i18n/ko-KR/permission-keys.json +++ b/web/i18n/ko-KR/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "API 확장 구성 관리", "app.access_config": "앱 접근 권한 구성", "app.acl.access_config": "접근 권한 보기 및 관리", + "app.acl.access_point_manage": "액세스 지점 보기 및 관리", "app.acl.delete": "앱 삭제", "app.acl.deploy": "앱 배포", "app.acl.edit": "앱 정보 편집 및 앱 오케스트레이션", diff --git a/web/i18n/lo-LA/permission-keys.json b/web/i18n/lo-LA/permission-keys.json index bfba51cc7e2..e04ef1f59ec 100644 --- a/web/i18n/lo-LA/permission-keys.json +++ b/web/i18n/lo-LA/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "ຈັດການການຕັ້ງຄ່າ API extension", "app.access_config": "ຕັ້ງຄ່າສິດການເຂົ້າເຖິງແອັບ", "app.acl.access_config": "ເບິ່ງ ແລະ ຈັດການສິດການເຂົ້າເຖິງ", + "app.acl.access_point_manage": "ເບິ່ງ ແລະ ຈັດການຈຸດເຂົ້າເຖິງ", "app.acl.delete": "ລຶບແອັບ", "app.acl.deploy": "ຕິດຕັ້ງແອັບ", "app.acl.edit": "ແກ້ໄຂຂໍ້ມູນແອັບ ແລະ ຈັດການລະບົບແອັບ", diff --git a/web/i18n/nl-NL/permission-keys.json b/web/i18n/nl-NL/permission-keys.json index 94fdf9d182c..8266b38ae22 100644 --- a/web/i18n/nl-NL/permission-keys.json +++ b/web/i18n/nl-NL/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "API-extensieconfiguratie beheren", "app.access_config": "Toegangsrechten voor app configureren", "app.acl.access_config": "Toegangsrechten bekijken en beheren", + "app.acl.access_point_manage": "Toegangspunten bekijken en beheren", "app.acl.delete": "App verwijderen", "app.acl.deploy": "App implementeren", "app.acl.edit": "App-informatie bewerken en app orkestreren", diff --git a/web/i18n/pl-PL/permission-keys.json b/web/i18n/pl-PL/permission-keys.json index 55619735f16..392f470914e 100644 --- a/web/i18n/pl-PL/permission-keys.json +++ b/web/i18n/pl-PL/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Zarządzaj konfiguracją rozszerzenia API", "app.access_config": "Konfiguruj uprawnienia dostępu do aplikacji", "app.acl.access_config": "Wyświetlaj uprawnienia dostępu i zarządzaj nimi", + "app.acl.access_point_manage": "Wyświetlaj punkty dostępu i zarządzaj nimi", "app.acl.delete": "Usuń aplikację", "app.acl.deploy": "Wdróż aplikację", "app.acl.edit": "Edytuj informacje o aplikacji i orkiestruj aplikację", diff --git a/web/i18n/pt-BR/permission-keys.json b/web/i18n/pt-BR/permission-keys.json index 32dda95b517..36dd03d5719 100644 --- a/web/i18n/pt-BR/permission-keys.json +++ b/web/i18n/pt-BR/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Gerenciar configuração de extensão de API", "app.access_config": "Configurar permissões de acesso ao aplicativo", "app.acl.access_config": "Visualizar e gerenciar permissões de acesso", + "app.acl.access_point_manage": "Visualizar e gerenciar pontos de acesso", "app.acl.delete": "Excluir aplicativo", "app.acl.deploy": "Implantar aplicativo", "app.acl.edit": "Editar informações e orquestrar o aplicativo", diff --git a/web/i18n/ro-RO/permission-keys.json b/web/i18n/ro-RO/permission-keys.json index 2610e185492..73225f7cd60 100644 --- a/web/i18n/ro-RO/permission-keys.json +++ b/web/i18n/ro-RO/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Gestionează configurația extensiei API", "app.access_config": "Configurează permisiunile de acces ale aplicației", "app.acl.access_config": "Vizualizează și gestionează permisiunile de acces", + "app.acl.access_point_manage": "Vizualizează și gestionează punctele de acces", "app.acl.delete": "Șterge aplicația", "app.acl.deploy": "Implementează aplicația", "app.acl.edit": "Editează informațiile aplicației și orchestrează aplicația", diff --git a/web/i18n/ru-RU/permission-keys.json b/web/i18n/ru-RU/permission-keys.json index 574f0e96add..bd986d65e95 100644 --- a/web/i18n/ru-RU/permission-keys.json +++ b/web/i18n/ru-RU/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Управление конфигурацией API-расширений", "app.access_config": "Настройка прав доступа к приложению", "app.acl.access_config": "Просмотр и управление правами доступа", + "app.acl.access_point_manage": "Просмотр и управление точками доступа", "app.acl.delete": "Удаление приложения", "app.acl.deploy": "Развертывание приложения", "app.acl.edit": "Редактирование информации о приложении и оркестрация приложения", diff --git a/web/i18n/sl-SI/permission-keys.json b/web/i18n/sl-SI/permission-keys.json index 544c6f92a8d..4a49494a2c2 100644 --- a/web/i18n/sl-SI/permission-keys.json +++ b/web/i18n/sl-SI/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Upravljanje konfiguracije razširitve API", "app.access_config": "Konfiguracija dovoljenj za dostop do aplikacije", "app.acl.access_config": "Ogled in upravljanje dovoljenj za dostop", + "app.acl.access_point_manage": "Ogled in upravljanje dostopnih točk", "app.acl.delete": "Izbriši aplikacijo", "app.acl.deploy": "Uvedi aplikacijo", "app.acl.edit": "Uredi podatke o aplikaciji in orkestriraj aplikacijo", diff --git a/web/i18n/th-TH/permission-keys.json b/web/i18n/th-TH/permission-keys.json index b7b9854abf0..0ad4047ef26 100644 --- a/web/i18n/th-TH/permission-keys.json +++ b/web/i18n/th-TH/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "จัดการการกําหนดค่าส่วนขยาย API", "app.access_config": "กําหนดค่าสิทธิ์การเข้าถึงแอป", "app.acl.access_config": "ดูและจัดการสิทธิ์การเข้าถึง", + "app.acl.access_point_manage": "ดูและจัดการจุดเข้าถึง", "app.acl.delete": "ลบแอป", "app.acl.deploy": "ปรับใช้แอป", "app.acl.edit": "แก้ไขข้อมูลแอปและจัดวางแอป", diff --git a/web/i18n/tr-TR/permission-keys.json b/web/i18n/tr-TR/permission-keys.json index 36ba8ec9709..781d9d00d38 100644 --- a/web/i18n/tr-TR/permission-keys.json +++ b/web/i18n/tr-TR/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "API uzantısı yapılandırmasını yönet", "app.access_config": "Uygulama erişim izinlerini yapılandır", "app.acl.access_config": "Erişim izinlerini görüntüle ve yönet", + "app.acl.access_point_manage": "Erişim noktalarını görüntüle ve yönet", "app.acl.delete": "Uygulamayı sil", "app.acl.deploy": "Uygulamayı dağıt", "app.acl.edit": "Uygulama bilgilerini düzenle ve uygulamayı orkestre et", diff --git a/web/i18n/uk-UA/permission-keys.json b/web/i18n/uk-UA/permission-keys.json index 861c83a4367..8cfd28d2548 100644 --- a/web/i18n/uk-UA/permission-keys.json +++ b/web/i18n/uk-UA/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Керування конфігурацією розширення API", "app.access_config": "Налаштування дозволів доступу до застосунку", "app.acl.access_config": "Переглядати дозволи доступу та керувати ними", + "app.acl.access_point_manage": "Переглядати точки доступу та керувати ними", "app.acl.delete": "Видалити застосунок", "app.acl.deploy": "Розгорнути застосунок", "app.acl.edit": "Редагувати інформацію про застосунок та оркеструвати застосунок", diff --git a/web/i18n/vi-VN/permission-keys.json b/web/i18n/vi-VN/permission-keys.json index 1e6662a9304..2290d6362e0 100644 --- a/web/i18n/vi-VN/permission-keys.json +++ b/web/i18n/vi-VN/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "Quản lý cấu hình phần mở rộng API", "app.access_config": "Cấu hình quyền truy cập ứng dụng", "app.acl.access_config": "Xem và quản lý quyền truy cập", + "app.acl.access_point_manage": "Xem và quản lý điểm truy cập", "app.acl.delete": "Xóa ứng dụng", "app.acl.deploy": "Triển khai ứng dụng", "app.acl.edit": "Chỉnh sửa thông tin và điều phối ứng dụng", diff --git a/web/i18n/zh-Hans/permission-keys.json b/web/i18n/zh-Hans/permission-keys.json index 91d9afb2cc8..db8184ee533 100644 --- a/web/i18n/zh-Hans/permission-keys.json +++ b/web/i18n/zh-Hans/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "管理API扩展", "app.access_config": "配置应用访问权限", "app.acl.access_config": "查看与管理访问权限", + "app.acl.access_point_manage": "查看与管理访问点", "app.acl.delete": "删除应用", "app.acl.deploy": "部署应用", "app.acl.edit": "编辑应用信息与编排应用", diff --git a/web/i18n/zh-Hant/permission-keys.json b/web/i18n/zh-Hant/permission-keys.json index 43350a59ada..8a74b37f086 100644 --- a/web/i18n/zh-Hant/permission-keys.json +++ b/web/i18n/zh-Hant/permission-keys.json @@ -3,6 +3,7 @@ "api_extension.manage": "管理API擴充配置", "app.access_config": "配置應用訪問權限", "app.acl.access_config": "檢視與管理存取權限", + "app.acl.access_point_manage": "檢視與管理存取點", "app.acl.delete": "刪除應用", "app.acl.deploy": "部署應用", "app.acl.edit": "編輯應用資訊與編排應用", diff --git a/web/utils/app-redirection.spec.ts b/web/utils/app-redirection.spec.ts index d736ed5a428..521a500652d 100644 --- a/web/utils/app-redirection.spec.ts +++ b/web/utils/app-redirection.spec.ts @@ -14,12 +14,22 @@ describe('app-redirection', () => { * - App mode (workflow, advanced-chat, chat, completion, agent-chat) */ describe('getRedirectionPath', () => { - it('returns access point path when app ACL cannot access guarded pages', () => { - const app = { id: 'app-123', mode: AppModeEnum.CHAT, permission_keys: [] } + it('returns access point path when app access point permission is granted', () => { + const app = { + id: 'app-123', + mode: AppModeEnum.CHAT, + permission_keys: [AppACLPermission.AccessPoint], + } const result = getRedirectionPath(app) expect(result).toBe('/app/app-123/access-point') }) + it('returns apps list path when app ACL cannot access guarded pages or access point', () => { + const app = { id: 'app-123', mode: AppModeEnum.CHAT, permission_keys: [] } + const result = getRedirectionPath(app) + expect(result).toBe('/apps') + }) + it('returns workflow path for workflow mode when app ACL can access layout', () => { const app = { id: 'app-123', @@ -92,7 +102,11 @@ describe('app-redirection', () => { }) it('handles different app IDs', () => { - const app1 = { id: 'abc-123', mode: AppModeEnum.CHAT, permission_keys: [] } + const app1 = { + id: 'abc-123', + mode: AppModeEnum.CHAT, + permission_keys: [AppACLPermission.AccessPoint], + } const app2 = { id: 'xyz-789', mode: AppModeEnum.WORKFLOW, @@ -129,7 +143,7 @@ describe('app-redirection', () => { const app = { id: 'app-123', mode: AppModeEnum.CHAT, - permission_keys: [AppACLPermission.AccessConfig], + permission_keys: [AppACLPermission.AccessConfig, AppACLPermission.AccessPoint], } expect(getRedirectionPath(app, { isRbacEnabled: false })).toBe('/app/app-123/access-point') @@ -173,8 +187,12 @@ describe('app-redirection', () => { /** * Tests that the redirection function is called with the correct path */ - it('calls redirection function with access point path when app ACL cannot access guarded pages', () => { - const app = { id: 'app-123', mode: AppModeEnum.CHAT, permission_keys: [] } + it('calls redirection function with access point path when access point permission is granted', () => { + const app = { + id: 'app-123', + mode: AppModeEnum.CHAT, + permission_keys: [AppACLPermission.AccessPoint], + } const mockRedirect = vi.fn() getRedirection(app, mockRedirect) diff --git a/web/utils/app-redirection.ts b/web/utils/app-redirection.ts index 4bdfce9a266..d73a350b3f3 100644 --- a/web/utils/app-redirection.ts +++ b/web/utils/app-redirection.ts @@ -34,7 +34,9 @@ export const getRedirectionPath = ( if (app.mode === AppModeEnum.WORKFLOW && appACLCapabilities.canDeploy) return `/app/${app.id}/deploy` - return `/app/${app.id}/access-point` + if (appACLCapabilities.canAccessPoint) return `/app/${app.id}/access-point` + + return '/apps' } export const getRedirection = ( diff --git a/web/utils/permission.spec.ts b/web/utils/permission.spec.ts index 75dff3e7afe..487c9c891b1 100644 --- a/web/utils/permission.spec.ts +++ b/web/utils/permission.spec.ts @@ -50,6 +50,15 @@ describe('permission', () => { expect(releaseCapabilities.canDeploy).toBe(false) }) + it('keeps access point permission independent from other app ACL permissions', () => { + const accessPointCapabilities = getAppACLCapabilities([AppACLPermission.AccessPoint]) + const layoutCapabilities = getAppACLCapabilities([AppACLPermission.ViewLayout]) + + expect(accessPointCapabilities.canAccessPoint).toBe(true) + expect(accessPointCapabilities.canAccessLayout).toBe(false) + expect(layoutCapabilities.canAccessPoint).toBe(false) + }) + it('keeps monitor, tracing config, and log/annotation permissions independent', () => { const monitorCapabilities = getAppACLCapabilities([AppACLPermission.Monitor]) const tracingCapabilities = getAppACLCapabilities([AppACLPermission.TracingConfig]) @@ -109,6 +118,7 @@ describe('permission', () => { }) expect(capabilities.canViewLayout).toBe(true) + expect(capabilities.canAccessPoint).toBe(true) expect(capabilities.canTestAndRun).toBe(true) expect(capabilities.canEdit).toBe(true) expect(capabilities.canImportExportDSL).toBe(true) diff --git a/web/utils/permission.ts b/web/utils/permission.ts index 063878b6607..30b0b3e0cbe 100644 --- a/web/utils/permission.ts +++ b/web/utils/permission.ts @@ -2,6 +2,7 @@ import type { PermissionKey } from '@/models/access-control' export const AppACLPermission = { Preview: 'app.acl.preview', + AccessPoint: 'app.acl.access_point_manage', ViewLayout: 'app.acl.view_layout', TestAndRun: 'app.acl.test_and_run', Edit: 'app.acl.edit', @@ -38,6 +39,7 @@ export type ResourceMaintainerPermissionOptions = { } type AppACLCapabilities = { + canAccessPoint: boolean canViewLayout: boolean canTestAndRun: boolean canEdit: boolean @@ -135,6 +137,11 @@ export const getAppACLCapabilities = ( ) return { + canAccessPoint: hasResourcePermission( + permissionKeys, + AppACLPermission.AccessPoint, + hasMaintainerPermissions, + ), canViewLayout, canTestAndRun, canEdit, From 9e07c2dac6fb454f0e78227f109d73bbf543898d Mon Sep 17 00:00:00 2001 From: "Byron.wang" Date: Mon, 31 Aug 2026 03:22:01 +0000 Subject: [PATCH 05/21] refactor(api): extract email registration application service (#41108) Co-authored-by: hjlarry --- api/.importlinter | 1 + .../console/auth/email_register.py | 223 +++---- api/controllers/console/flask_admission.py | 16 + api/controllers/console/wraps.py | 44 +- api/extensions/ext_application_services.py | 34 ++ api/repositories/account_repository.py | 8 + .../account_email_registration_adapters.py | 230 +++++++ .../account_email_registration_service.py | 207 +++++++ api/services/account_errors.py | 40 ++ api/services/account_ports.py | 2 + api/services/account_service.py | 86 --- api/services/entities/account_entities.py | 24 + .../console/auth/test_email_register.py | 561 ++++++++---------- .../auth/test_email_register_language.py | 44 -- .../controllers/console/test_wraps.py | 58 ++ .../test_ext_application_services.py | 13 + .../repositories/test_account_repository.py | 14 + ...est_account_email_registration_adapters.py | 176 ++++++ ...test_account_email_registration_service.py | 267 +++++++++ 19 files changed, 1445 insertions(+), 603 deletions(-) create mode 100644 api/services/account_email_registration_adapters.py create mode 100644 api/services/account_email_registration_service.py delete mode 100644 api/tests/unit_tests/controllers/console/auth/test_email_register_language.py create mode 100644 api/tests/unit_tests/services/test_account_email_registration_adapters.py create mode 100644 api/tests/unit_tests/services/test_account_email_registration_service.py diff --git a/api/.importlinter b/api/.importlinter index f3609e5826a..7cf69c0c515 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -207,6 +207,7 @@ source_modules = services.account_avatar_service services.account_change_email_ports services.account_change_email_service + services.account_email_registration_service services.account_deletion_service services.account_deletion_feedback_service services.account_education_service diff --git a/api/controllers/console/auth/email_register.py b/api/controllers/console/auth/email_register.py index e81acbed99d..e0bfddccfa0 100644 --- a/api/controllers/console/auth/email_register.py +++ b/api/controllers/console/auth/email_register.py @@ -2,8 +2,6 @@ from flask import request from flask_restx import Resource from pydantic import BaseModel, Field, field_validator -from configs import dify_config -from constants.languages import get_valid_language, languages from controllers.common.fields import SimpleResultDataResponse, VerificationTokenResponse from controllers.common.schema import register_response_schema_models, register_schema_models from controllers.console import console_ns @@ -11,31 +9,35 @@ from controllers.console.auth.error import ( EmailAlreadyInUseError, EmailCodeError, EmailRegisterLimitError, + EmailRegisterRateLimitExceededError, InvalidEmailError, InvalidTokenError, NormalizedEmailAlreadyInUseError, PasswordMismatchError, ) -from enums import DeploymentEdition -from extensions.ext_database import db +from controllers.console.flask_admission import console_email_registration_admission +from controllers.console.wraps import model_validate +from extensions.ext_application_services import application_services from fields.base import ResponseModel -from libs.helper import EmailStr, extract_remote_ip +from libs.helper import EmailStr, dump_response, extract_remote_ip from libs.helper import timezone as validate_timezone_string 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 ( +from services.account_errors import ( + AccountEmailAlreadyInUseError, + AccountEmailDomainSuspendedError, + AccountEmailFrozenError, AccountNormalizedEmailAlreadyInUseError, - AccountRegisterError, - SeatsLimitExceededError, -) -from services.errors.account import ( - EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, + EmailRegistrationPasswordMismatchError, + EmailRegistrationSeatsLimitError, + EmailRegistrationSendIPLimitedError, + EmailRegistrationSendRateLimitError, + EmailRegistrationVerificationLimitError, + InvalidEmailRegistrationAddressError, + InvalidEmailRegistrationCodeError, + InvalidEmailRegistrationTokenError, ) from ..error import AccountInFreezeError, EmailDomainSuspendedError, EmailSendIpLimitError, SeatsLimitExceeded -from ..wraps import email_password_login_enabled, email_register_enabled, model_validate, setup_required class EmailRegisterSendPayload(BaseModel): @@ -91,146 +93,91 @@ register_response_schema_models( @console_ns.route("/email-register/send-email") class EmailRegisterSendEmailApi(Resource): - @setup_required - @email_password_login_enabled - @email_register_enabled @console_ns.expect(console_ns.models[EmailRegisterSendPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultDataResponse.__name__]) + @console_email_registration_admission @model_validate(EmailRegisterSendPayload) - def post(self, req_data: EmailRegisterSendPayload): - normalized_email = req_data.email.lower() - - ip_address = extract_remote_ip(request) - if AccountService.is_email_send_ip_limit(ip_address): - raise EmailSendIpLimitError() - language = "en-US" - if req_data.language is not None and req_data.language in languages: - language = req_data.language - - 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) - return {"result": "success", "data": token} + def post(self, args: EmailRegisterSendPayload): + try: + token = application_services().accounts.email_registration.send_code( + remote_ip=extract_remote_ip(request), + requested_email=args.email, + requested_language=args.language, + ) + except EmailRegistrationSendIPLimitedError: + raise EmailSendIpLimitError() from None + except EmailRegistrationSendRateLimitError as error: + raise EmailRegisterRateLimitExceededError(error.retry_after_minutes) from None + except AccountEmailDomainSuspendedError: + raise EmailDomainSuspendedError() from None + except AccountEmailFrozenError: + raise AccountInFreezeError() from None + return dump_response(SimpleResultDataResponse, {"result": "success", "data": token}) @console_ns.route("/email-register/validity") class EmailRegisterCheckApi(Resource): - @setup_required - @email_password_login_enabled - @email_register_enabled @console_ns.expect(console_ns.models[EmailRegisterValidityPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[VerificationTokenResponse.__name__]) + @console_email_registration_admission @model_validate(EmailRegisterValidityPayload) - def post(self, req_data: EmailRegisterValidityPayload): - - user_email = req_data.email.lower() - - is_email_register_error_rate_limit = AccountService.is_email_register_error_rate_limit(user_email) - if is_email_register_error_rate_limit: - raise EmailRegisterLimitError() - - token_data = AccountService.get_email_register_data(req_data.token) - if token_data is None: - raise InvalidTokenError() - - token_email = token_data.get("email") - normalized_token_email = token_email.lower() if isinstance(token_email, str) else token_email - - if user_email != normalized_token_email: - raise InvalidEmailError() - - if req_data.code != token_data.get("code"): - AccountService.add_email_register_error_rate_limit(user_email) - raise EmailCodeError() - - # Verified, revoke the first token - AccountService.revoke_email_register_token(req_data.token) - - # Refresh token data by generating a new token - _, new_token = AccountService.generate_email_register_token( - user_email, code=req_data.code, additional_data={"phase": "register"} + def post(self, args: EmailRegisterValidityPayload): + try: + verification = application_services().accounts.email_registration.verify_code( + email=args.email, + code=args.code, + token=args.token, + ) + except EmailRegistrationVerificationLimitError: + raise EmailRegisterLimitError() from None + except InvalidEmailRegistrationTokenError: + raise InvalidTokenError() from None + except InvalidEmailRegistrationAddressError: + raise InvalidEmailError() from None + except InvalidEmailRegistrationCodeError: + raise EmailCodeError() from None + return dump_response( + VerificationTokenResponse, + { + "is_valid": True, + "email": verification.email, + "token": verification.token, + }, ) - AccountService.reset_email_register_error_rate_limit(user_email) - return {"is_valid": True, "email": normalized_token_email, "token": new_token} - @console_ns.route("/email-register") class EmailRegisterResetApi(Resource): - @setup_required - @email_password_login_enabled - @email_register_enabled @console_ns.expect(console_ns.models[EmailRegisterResetPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[EmailRegisterResetResponse.__name__]) + @console_email_registration_admission @model_validate(EmailRegisterResetPayload) - def post(self, req_data: EmailRegisterResetPayload): - - # Validate passwords match - if req_data.new_password != req_data.password_confirm: - raise PasswordMismatchError() - - # Validate token and get register data - register_data = AccountService.get_email_register_data(req_data.token) - if not register_data: - raise InvalidTokenError() - # Must use token in reset phase - if register_data.get("phase", "") != "register": - raise InvalidTokenError() - - # Revoke token to prevent reuse - AccountService.revoke_email_register_token(req_data.token) - - email = register_data.get("email", "") - normalized_email = email.lower() - - account = AccountService.get_account_by_email_with_case_fallback(email, session=db.session()) - - if account: - raise EmailAlreadyInUseError() - - ip_address = extract_remote_ip(request) - account = self._create_new_account( - email=normalized_email, - password=req_data.password_confirm, - timezone=req_data.timezone, - language=req_data.language, - ip_address=ip_address, - ) - token_pair = AccountService.login(account=account, session=db.session(), ip_address=ip_address) - AccountService.reset_login_error_rate_limit(normalized_email) - - return {"result": "success", "data": token_pair.model_dump()} - - def _create_new_account( - self, - email: str, - password: str, - timezone: str | None = None, - language: str | None = None, - ip_address: str | None = None, - ) -> Account: + def post(self, args: EmailRegisterResetPayload): try: - return AccountService.create_account_and_tenant( - email=email, - name=email, - password=password, - interface_language=get_valid_language(language), - timezone=timezone, - ip_address=ip_address, - check_normalized_email=True, - session=db.session(), + token_pair = application_services().accounts.email_registration.register( + remote_ip=extract_remote_ip(request), + token=args.token, + new_password=args.new_password, + password_confirm=args.password_confirm, + language=args.language, + timezone=args.timezone, ) - except SeatsLimitExceededError: - raise SeatsLimitExceeded() - except EmailDomainSuspendedRegistrationError as exc: - raise EmailDomainSuspendedError() from exc - except AccountNormalizedEmailAlreadyInUseError as exc: - raise NormalizedEmailAlreadyInUseError() from exc - except AccountRegisterError as exc: - raise AccountInFreezeError() from exc + except EmailRegistrationPasswordMismatchError: + raise PasswordMismatchError() from None + except InvalidEmailRegistrationTokenError: + raise InvalidTokenError() from None + except AccountNormalizedEmailAlreadyInUseError: + raise NormalizedEmailAlreadyInUseError() from None + except AccountEmailAlreadyInUseError: + raise EmailAlreadyInUseError() from None + except EmailRegistrationSeatsLimitError: + raise SeatsLimitExceeded() from None + except AccountEmailDomainSuspendedError: + raise EmailDomainSuspendedError() from None + except AccountEmailFrozenError: + raise AccountInFreezeError() from None + + return dump_response( + EmailRegisterResetResponse, + {"result": "success", "data": token_pair}, + ) diff --git a/api/controllers/console/flask_admission.py b/api/controllers/console/flask_admission.py index 5eafc9c4741..eb300128aed 100644 --- a/api/controllers/console/flask_admission.py +++ b/api/controllers/console/flask_admission.py @@ -22,6 +22,22 @@ from libs.login import current_account_with_tenant, login_required from machinery.context import RequestContext from machinery.errors import AdmissionConfigurationError from models.account import TenantAccountRole +from services.feature_service import FeatureService + + +def console_email_registration_admission[T, **P, R]( + view: Callable[Concatenate[T, P], R], +) -> Callable[Concatenate[T, P], R | Response]: + """Apply the complete admission policy for anonymous email registration.""" + + @wraps(view) + def check_registration_features(self: T, /, *args: P.args, **kwargs: P.kwargs) -> R: + features = FeatureService.get_system_features() + if not features.enable_email_password_login or not features.is_allow_register: + abort(403) + return view(self, *args, **kwargs) + + return setup_required(check_registration_features) def console_account_admission[T, **P, R]( diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index 7850f4d206a..da1d5584af6 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -352,19 +352,6 @@ def email_password_login_enabled[**P, R](view: Callable[P, R]) -> Callable[P, R] return decorated -def email_register_enabled[**P, R](view: Callable[P, R]) -> Callable[P, R]: - @wraps(view) - def decorated(*args: P.args, **kwargs: P.kwargs): - features = FeatureService.get_system_features() - if features.is_allow_register: - return view(*args, **kwargs) - - # otherwise, return 403 - abort(403) - - return decorated - - def enable_change_email[**P, R](view: Callable[P, R]) -> Callable[P, R]: @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): @@ -652,6 +639,23 @@ def with_current_user_id[T, **P, R]( return decorated +def validate_request[M: BaseModel](model: type[M]) -> M: + """Parse and validate the current request without exposing submitted values.""" + + if request.method == "GET": + raw = request.args.to_dict(flat=True) + elif request.method == "DELETE": + raw = request.args.to_dict(flat=True) or (request.get_json(silent=True) or {}) + else: + raw = request.get_json(silent=True) or {} + + try: + return model.model_validate(raw) + except ValidationError as exc: + errors = exc.errors(include_url=False, include_input=False, include_context=False) + raise UnprocessableEntity(json.dumps(errors)) from None + + def model_validate[T, M: BaseModel, **P, R]( model: type[M], ) -> Callable[ @@ -671,19 +675,7 @@ def model_validate[T, M: BaseModel, **P, R]( ) -> Callable[Concatenate[T, P], R]: @wraps(view) def wrapper(self: T, *args: P.args, **kwargs: P.kwargs) -> R: - if request.method == "GET": - raw = request.args.to_dict(flat=True) - elif request.method == "DELETE": - raw = request.args.to_dict(flat=True) or (request.get_json(silent=True) or {}) - else: - raw = request.get_json(silent=True) or {} - - try: - validated = model.model_validate(raw) - except ValidationError as exc: - raise UnprocessableEntity(exc.json()) - - return view(self, validated, *args, **kwargs) + return view(self, validate_request(model), *args, **kwargs) return wrapper diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index e29b92c8617..7a482248cce 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -70,6 +70,16 @@ from services.account_deletion_adapters import ( from services.account_deletion_feedback_service import AccountDeletionFeedbackService from services.account_deletion_service import AccountDeletionService from services.account_education_service import AccountEducationService +from services.account_email_registration_adapters import ( + AccountServiceRegistrationGateway, + BillingAccountRegistrationPolicyGateway, + CeleryEmailRegistrationNotificationGateway, + RateLimiterEmailRegistrationSendLimiter, + RedisEmailRegistrationSecurityGateway, + SecureEmailRegistrationCodeGenerator, + TokenManagerEmailRegistrationTokenGateway, +) +from services.account_email_registration_service import AccountEmailRegistrationService from services.account_initialization_service import AccountInitializationService from services.account_integration_service import AccountIntegrationService from services.account_password_hasher import LegacyAccountPasswordHasher @@ -150,6 +160,7 @@ def _is_user_allowed_to_access_webapp(user_id: str, app_id: str) -> bool: class AccountServices: avatar: AccountAvatarService change_email: AccountChangeEmailService + email_registration: AccountEmailRegistrationService deletion: AccountDeletionService deletion_feedback: AccountDeletionFeedbackService education: AccountEducationService @@ -278,6 +289,29 @@ def build_application_services( billing_enabled=deployment_edition == DeploymentEdition.CLOUD, ), ), + email_registration=AccountEmailRegistrationService( + accounts=accounts, + tokens=TokenManagerEmailRegistrationTokenGateway(), + codes=SecureEmailRegistrationCodeGenerator(), + notifications=CeleryEmailRegistrationNotificationGateway(), + send_limits=RateLimiterEmailRegistrationSendLimiter( + rate_limiter=RateLimiter( + prefix="email_register_rate_limit", + max_attempts=1, + time_window=60, + redis_client=redis, + ) + ), + security=RedisEmailRegistrationSecurityGateway( + redis=redis, + verification_failure_limit=5, + verification_lockout_duration=dify_config.EMAIL_REGISTER_LOCKOUT_DURATION, + ), + account_policy=BillingAccountRegistrationPolicyGateway( + enabled=deployment_edition == DeploymentEdition.CLOUD, + ), + registration=AccountServiceRegistrationGateway(session_factory=database_client), + ), deletion=AccountDeletionService( accounts=accounts, memberships=workspace_query_repository, diff --git a/api/repositories/account_repository.py b/api/repositories/account_repository.py index 1e7fb55e0d3..54568c7321c 100644 --- a/api/repositories/account_repository.py +++ b/api/repositories/account_repository.py @@ -31,6 +31,14 @@ class SQLAlchemyAccountRepository(AccountRepository): account = session.get(Account, account_id) return self._to_snapshot(account) if account is not None else None + @override + def find_by_email(self, email: str) -> AccountSnapshot | None: + with self._session_factory() as session: + account = session.scalar(select(Account).where(Account.email == email).limit(1)) + if account is None and email != email.lower(): + account = session.scalar(select(Account).where(Account.email == email.lower()).limit(1)) + return self._to_snapshot(account) if account is not None else None + @override def get_credentials(self, account_id: str) -> AccountCredentials | None: with self._session_factory() as session: diff --git a/api/services/account_email_registration_adapters.py b/api/services/account_email_registration_adapters.py new file mode 100644 index 00000000000..5bd25b33b6e --- /dev/null +++ b/api/services/account_email_registration_adapters.py @@ -0,0 +1,230 @@ +"""Infrastructure adapters for account email registration.""" + +import logging +import secrets +from typing import override + +from redis import RedisError +from sqlalchemy.orm import Session, sessionmaker + +from extensions.ext_redis import RedisClientWrapper +from libs.helper import RateLimiter, TokenManager +from models.account import Account +from services.account_email_registration_service import ( + AccountRegistrationGateway, + AccountRegistrationPolicyGateway, + EmailRegistrationCodeGenerator, + EmailRegistrationNotificationGateway, + EmailRegistrationSecurityGateway, + EmailRegistrationSendLimiter, + EmailRegistrationTokenGateway, +) +from services.account_errors import ( + AccountEmailDomainSuspendedError, + AccountEmailFrozenError, + AccountNormalizedEmailAlreadyInUseError, + EmailRegistrationSeatsLimitError, +) +from services.account_service import AccountService +from services.billing_service import BillingService +from services.entities.account_entities import ( + AccountEmailRegistrationPhase, + AccountEmailRegistrationToken, + AccountSessionTokens, +) +from services.errors.account import ( + AccountNormalizedEmailAlreadyInUseError as AccountNormalizedEmailAlreadyInUseServiceError, +) +from services.errors.account import AccountRegisterError, EmailDomainSuspendedError, SeatsLimitExceededError +from tasks.mail_register_task import send_email_register_mail_task, send_email_register_mail_task_when_account_exist + +logger = logging.getLogger(__name__) + + +class TokenManagerEmailRegistrationTokenGateway(EmailRegistrationTokenGateway): + @override + def get(self, token: str) -> AccountEmailRegistrationToken | None: + payload = TokenManager.get_token_data(token, "email_register") + if payload is None: + return None + email = payload.get("email") + code = payload.get("code") + phase_value = payload.get("phase") + if not isinstance(email, str) or not isinstance(code, str): + return None + if phase_value is None: + phase = None + else: + try: + phase = AccountEmailRegistrationPhase(phase_value) + except (TypeError, ValueError): + return None + return AccountEmailRegistrationToken(email=email, code=code, phase=phase) + + @override + def issue(self, token_data: AccountEmailRegistrationToken) -> str: + additional_data = {"code": token_data.code} + if token_data.phase is not None: + additional_data["phase"] = token_data.phase.value + return TokenManager.generate_token( + email=token_data.email, + token_type="email_register", + additional_data=additional_data, + ) + + @override + def revoke(self, token: str) -> None: + TokenManager.revoke_token(token, "email_register") + + +class SecureEmailRegistrationCodeGenerator(EmailRegistrationCodeGenerator): + @override + def generate(self) -> str: + return "".join(str(secrets.randbelow(exclusive_upper_bound=10)) for _ in range(6)) + + +class CeleryEmailRegistrationNotificationGateway(EmailRegistrationNotificationGateway): + @override + def send_code(self, *, email: str, code: str, language: str) -> None: + send_email_register_mail_task.delay(language=language, to=email, code=code) + + @override + def send_account_exists(self, *, email: str, account_name: str, language: str) -> None: + send_email_register_mail_task_when_account_exist.delay( + language=language, + to=email, + account_name=account_name, + ) + + +class RateLimiterEmailRegistrationSendLimiter(EmailRegistrationSendLimiter): + def __init__(self, *, rate_limiter: RateLimiter) -> None: + self._rate_limiter = rate_limiter + + @override + def is_limited(self, email: str) -> bool: + return self._rate_limiter.is_rate_limited(email) + + @override + def record(self, email: str) -> None: + self._rate_limiter.increment_rate_limit(email) + + @property + @override + def retry_after_minutes(self) -> int: + return int(self._rate_limiter.time_window / 60) + + +class RedisEmailRegistrationSecurityGateway(EmailRegistrationSecurityGateway): + def __init__( + self, + *, + redis: RedisClientWrapper, + verification_failure_limit: int, + verification_lockout_duration: int, + ) -> None: + self._redis = redis + self._verification_failure_limit = verification_failure_limit + self._verification_lockout_duration = verification_lockout_duration + + @override + def is_ip_limited(self, ip_address: str) -> bool: + return AccountService.is_email_send_ip_limit(ip_address) is True + + @override + def is_verification_limited(self, email: str) -> bool: + try: + count = self._redis.get(self._verification_key(email)) + return count is not None and int(count) > self._verification_failure_limit + except RedisError: + logger.warning("Failed to read email-registration verification limit", exc_info=True) + return False + + @override + def record_verification_failure(self, email: str) -> None: + try: + key = self._verification_key(email) + count = int(self._redis.get(key) or 0) + 1 + self._redis.setex(key, self._verification_lockout_duration, count) + except RedisError: + logger.warning("Failed to record email-registration verification failure", exc_info=True) + return None + + @override + def reset_verification_failures(self, email: str) -> None: + try: + self._redis.delete(self._verification_key(email)) + except RedisError: + logger.warning("Failed to reset email-registration verification failures", exc_info=True) + return None + + @override + def reset_login_failures(self, email: str) -> None: + AccountService.reset_login_error_rate_limit(email) + + @staticmethod + def _verification_key(email: str) -> str: + return f"email_register_error_rate_limit:{email}" + + +class BillingAccountRegistrationPolicyGateway(AccountRegistrationPolicyGateway): + def __init__(self, *, enabled: bool) -> None: + self._enabled = enabled + + @override + def get_freeze_type(self, email: str) -> str | None: + if not self._enabled: + return None + return BillingService.get_email_freeze_type(email) + + +class AccountServiceRegistrationGateway(AccountRegistrationGateway): + """Compatibility adapter around account provisioning and login internals.""" + + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def create( + self, + *, + email: str, + password: str, + interface_language: str, + timezone: str | None, + ip_address: str, + ) -> str: + with self._session_factory() as session: + try: + account = AccountService.create_account_and_tenant( + email=email, + name=email, + password=password, + interface_language=interface_language, + timezone=timezone, + ip_address=ip_address, + check_normalized_email=True, + session=session, + ) + except SeatsLimitExceededError as exc: + raise EmailRegistrationSeatsLimitError from exc + except EmailDomainSuspendedError as exc: + raise AccountEmailDomainSuspendedError from exc + except AccountNormalizedEmailAlreadyInUseServiceError as exc: + raise AccountNormalizedEmailAlreadyInUseError from exc + except AccountRegisterError as exc: + raise AccountEmailFrozenError from exc + return account.id + + @override + def login(self, account_id: str, *, ip_address: str) -> AccountSessionTokens: + with self._session_factory() as session: + account = session.get(Account, account_id) + if account is None: + raise RuntimeError("newly registered account no longer exists") + token_pair = AccountService.login(account=account, session=session, ip_address=ip_address) + return AccountSessionTokens( + access_token=token_pair.access_token, + refresh_token=token_pair.refresh_token, + csrf_token=token_pair.csrf_token, + ) diff --git a/api/services/account_email_registration_service.py b/api/services/account_email_registration_service.py new file mode 100644 index 00000000000..2379f220254 --- /dev/null +++ b/api/services/account_email_registration_service.py @@ -0,0 +1,207 @@ +"""Application service for the account email-registration use case.""" + +from typing import Protocol + +from constants.languages import get_valid_language, languages +from services.account_errors import ( + AccountEmailAlreadyInUseError, + AccountEmailDomainSuspendedError, + AccountEmailFrozenError, + EmailRegistrationPasswordMismatchError, + EmailRegistrationSendIPLimitedError, + EmailRegistrationSendRateLimitError, + EmailRegistrationVerificationLimitError, + InvalidEmailRegistrationAddressError, + InvalidEmailRegistrationCodeError, + InvalidEmailRegistrationTokenError, +) +from services.account_ports import AccountRepository +from services.entities.account_entities import ( + AccountEmailRegistrationPhase, + AccountEmailRegistrationToken, + AccountEmailRegistrationVerification, + AccountSessionTokens, +) + + +class EmailRegistrationTokenGateway(Protocol): + def get(self, token: str) -> AccountEmailRegistrationToken | None: ... + + def issue(self, token_data: AccountEmailRegistrationToken) -> str: ... + + def revoke(self, token: str) -> None: ... + + +class EmailRegistrationCodeGenerator(Protocol): + def generate(self) -> str: ... + + +class EmailRegistrationNotificationGateway(Protocol): + def send_code(self, *, email: str, code: str, language: str) -> None: ... + + def send_account_exists(self, *, email: str, account_name: str, language: str) -> None: ... + + +class EmailRegistrationSendLimiter(Protocol): + def is_limited(self, email: str) -> bool: ... + + def record(self, email: str) -> None: ... + + @property + def retry_after_minutes(self) -> int: ... + + +class EmailRegistrationSecurityGateway(Protocol): + def is_ip_limited(self, ip_address: str) -> bool: ... + + def is_verification_limited(self, email: str) -> bool: ... + + def record_verification_failure(self, email: str) -> None: ... + + def reset_verification_failures(self, email: str) -> None: ... + + def reset_login_failures(self, email: str) -> None: ... + + +class AccountRegistrationPolicyGateway(Protocol): + def get_freeze_type(self, email: str) -> str | None: ... + + +class AccountRegistrationGateway(Protocol): + def create( + self, + *, + email: str, + password: str, + interface_language: str, + timezone: str | None, + ip_address: str, + ) -> str: ... + + def login(self, account_id: str, *, ip_address: str) -> AccountSessionTokens: ... + + +class AccountEmailRegistrationService: + def __init__( + self, + *, + accounts: AccountRepository, + tokens: EmailRegistrationTokenGateway, + codes: EmailRegistrationCodeGenerator, + notifications: EmailRegistrationNotificationGateway, + send_limits: EmailRegistrationSendLimiter, + security: EmailRegistrationSecurityGateway, + account_policy: AccountRegistrationPolicyGateway, + registration: AccountRegistrationGateway, + ) -> None: + self._accounts = accounts + self._tokens = tokens + self._codes = codes + self._notifications = notifications + self._send_limits = send_limits + self._security = security + self._account_policy = account_policy + self._registration = registration + + def send_code( + self, + *, + remote_ip: str, + requested_email: str, + requested_language: str | None, + ) -> str: + if self._security.is_ip_limited(remote_ip): + raise EmailRegistrationSendIPLimitedError + + normalized_email = requested_email.lower() + self._ensure_email_allowed(normalized_email) + account = self._accounts.find_by_email(requested_email) + delivery_email = account.email if account is not None else normalized_email + if self._send_limits.is_limited(delivery_email): + raise EmailRegistrationSendRateLimitError(self._send_limits.retry_after_minutes) + + language = requested_language if requested_language is not None and requested_language in languages else "en-US" + code = self._codes.generate() + token = self._tokens.issue(AccountEmailRegistrationToken(email=delivery_email, code=code)) + if account is None: + self._notifications.send_code(email=delivery_email, code=code, language=language) + else: + self._notifications.send_account_exists( + email=delivery_email, + account_name=account.name, + language=language, + ) + self._send_limits.record(delivery_email) + return token + + def verify_code( + self, + *, + email: str, + code: str, + token: str, + ) -> AccountEmailRegistrationVerification: + normalized_email = email.lower() + if self._security.is_verification_limited(normalized_email): + raise EmailRegistrationVerificationLimitError + + token_data = self._tokens.get(token) + if token_data is None: + raise InvalidEmailRegistrationTokenError + normalized_token_email = token_data.email.lower() + if normalized_email != normalized_token_email: + raise InvalidEmailRegistrationAddressError + if code != token_data.code: + self._security.record_verification_failure(normalized_email) + raise InvalidEmailRegistrationCodeError + + self._tokens.revoke(token) + verified_token = self._tokens.issue( + AccountEmailRegistrationToken( + email=normalized_email, + code=code, + phase=AccountEmailRegistrationPhase.REGISTER, + ) + ) + self._security.reset_verification_failures(normalized_email) + return AccountEmailRegistrationVerification(email=normalized_token_email, token=verified_token) + + def register( + self, + *, + remote_ip: str, + token: str, + new_password: str, + password_confirm: str, + language: str | None, + timezone: str | None, + ) -> AccountSessionTokens: + if new_password != password_confirm: + raise EmailRegistrationPasswordMismatchError + + token_data = self._tokens.get(token) + if token_data is None or token_data.phase != AccountEmailRegistrationPhase.REGISTER: + raise InvalidEmailRegistrationTokenError + self._tokens.revoke(token) + + normalized_email = token_data.email.lower() + if self._accounts.find_by_email(token_data.email) is not None: + raise AccountEmailAlreadyInUseError + + account_id = self._registration.create( + email=normalized_email, + password=password_confirm, + interface_language=get_valid_language(language), + timezone=timezone, + ip_address=remote_ip, + ) + tokens = self._registration.login(account_id, ip_address=remote_ip) + self._security.reset_login_failures(normalized_email) + return tokens + + def _ensure_email_allowed(self, email: str) -> None: + freeze_type = self._account_policy.get_freeze_type(email) + if freeze_type == "email_domain_suspended": + raise AccountEmailDomainSuspendedError + if freeze_type: + raise AccountEmailFrozenError diff --git a/api/services/account_errors.py b/api/services/account_errors.py index c902c7d331b..115e0e6f511 100644 --- a/api/services/account_errors.py +++ b/api/services/account_errors.py @@ -85,6 +85,46 @@ class AccountEmailAlreadyInUseError(AccountApplicationError): """The target email already belongs to an account.""" +class AccountNormalizedEmailAlreadyInUseError(AccountEmailAlreadyInUseError): + """A normalized equivalent of the target email already belongs to an account.""" + + +class EmailRegistrationSendIPLimitedError(AccountApplicationError): + """The caller IP exceeded the registration-email send policy.""" + + +class EmailRegistrationSendRateLimitError(AccountApplicationError): + """Too many registration messages were requested for the address.""" + + def __init__(self, retry_after_minutes: int) -> None: + super().__init__(retry_after_minutes) + self.retry_after_minutes = retry_after_minutes + + +class EmailRegistrationVerificationLimitError(AccountApplicationError): + """Too many invalid registration-code attempts were made.""" + + +class InvalidEmailRegistrationTokenError(AccountApplicationError): + """The registration token is absent, malformed, or in the wrong phase.""" + + +class InvalidEmailRegistrationAddressError(AccountApplicationError): + """The request address does not match the registration token.""" + + +class InvalidEmailRegistrationCodeError(AccountApplicationError): + """The verification code does not match the registration token.""" + + +class EmailRegistrationPasswordMismatchError(AccountApplicationError): + """The registration password confirmation does not match.""" + + +class EmailRegistrationSeatsLimitError(AccountApplicationError): + """The deployment has no licensed seat available for another account.""" + + class EducationDiscountPausedError(AccountApplicationError): """Education discount activation is temporarily paused.""" diff --git a/api/services/account_ports.py b/api/services/account_ports.py index 39afd92bd1f..78792664127 100644 --- a/api/services/account_ports.py +++ b/api/services/account_ports.py @@ -19,6 +19,8 @@ from services.entities.account_entities import ( class AccountRepository(Protocol): def get(self, account_id: str) -> AccountSnapshot | None: ... + def find_by_email(self, email: str) -> AccountSnapshot | None: ... + def get_credentials(self, account_id: str) -> AccountCredentials | None: ... def update_profile(self, account_id: str, changes: AccountProfileChanges) -> AccountSnapshot | None: ... diff --git a/api/services/account_service.py b/api/services/account_service.py index 586619abe1d..b022187e14d 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -93,7 +93,6 @@ from tasks.mail_owner_transfer_task import ( send_old_owner_transfer_notify_email_task, send_owner_transfer_confirm_task, ) -from tasks.mail_register_task import send_email_register_mail_task, send_email_register_mail_task_when_account_exist from tasks.mail_reset_password_task import ( send_reset_password_mail_task, send_reset_password_mail_task_when_account_not_exist, @@ -157,7 +156,6 @@ class AccountService: CHANGE_EMAIL_PHASE_NEW = ChangeEmailPhase.NEW_EMAIL reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=1, time_window=60 * 1) - email_register_rate_limiter = RateLimiter(prefix="email_register_rate_limit", max_attempts=1, time_window=60 * 1) email_code_login_rate_limiter = RateLimiter( prefix="email_code_login_rate_limit", max_attempts=3, time_window=300 * 1 ) @@ -168,7 +166,6 @@ class AccountService: FORGOT_PASSWORD_MAX_ERROR_LIMITS = 5 CHANGE_EMAIL_MAX_ERROR_LIMITS = 5 OWNER_TRANSFER_MAX_ERROR_LIMITS = 5 - EMAIL_REGISTER_MAX_ERROR_LIMITS = 5 @staticmethod def _resolve_role_id_by_tag(tenant_id: str, account_id: str, tag: str) -> str: @@ -679,40 +676,6 @@ class AccountService: cls.reset_password_rate_limiter.increment_rate_limit(account_email) return token - @classmethod - def send_email_register_email( - cls, - account: Account | None = None, - email: str | None = None, - language: str = "en-US", - ): - account_email = account.email if account else email - if account_email is None: - raise ValueError("Email must be provided.") - - if cls.email_register_rate_limiter.is_rate_limited(account_email): - from controllers.console.auth.error import EmailRegisterRateLimitExceededError - - raise EmailRegisterRateLimitExceededError(int(cls.email_register_rate_limiter.time_window / 60)) - - code, token = cls.generate_email_register_token(account_email) - - if account: - send_email_register_mail_task_when_account_exist.delay( - language=language, - to=account_email, - account_name=account.name, - ) - - else: - send_email_register_mail_task.delay( - language=language, - to=account_email, - code=code, - ) - cls.email_register_rate_limiter.increment_rate_limit(account_email) - return token - @classmethod def send_change_email_email( cls, @@ -866,19 +829,6 @@ class AccountService: ) return code, token - @classmethod - def generate_email_register_token( - cls, - email: str, - code: str | None = None, - additional_data: dict[str, Any] = {}, - ): - if not code: - code = "".join([str(secrets.randbelow(exclusive_upper_bound=10)) for _ in range(6)]) - additional_data["code"] = code - token = TokenManager.generate_token(email=email, token_type="email_register", additional_data=additional_data) - return code, token - @classmethod def generate_change_email_token( cls, @@ -916,10 +866,6 @@ class AccountService: def revoke_reset_password_token(cls, token: str): TokenManager.revoke_token(token, "reset_password") - @classmethod - def revoke_email_register_token(cls, token: str): - TokenManager.revoke_token(token, "email_register") - @classmethod def revoke_change_email_token(cls, token: str): TokenManager.revoke_token(token, "change_email") @@ -932,10 +878,6 @@ class AccountService: def get_reset_password_data(cls, token: str) -> dict[str, Any] | None: return TokenManager.get_token_data(token, "reset_password") - @classmethod - def get_email_register_data(cls, token: str) -> dict[str, Any] | None: - return TokenManager.get_token_data(token, "email_register") - @classmethod def get_change_email_data(cls, token: str) -> ChangeEmailTokenData | None: token_data = TokenManager.get_token_data(token, "change_email") @@ -1066,16 +1008,6 @@ class AccountService: count = int(count) + 1 redis_client.setex(key, dify_config.FORGOT_PASSWORD_LOCKOUT_DURATION, count) - @staticmethod - @redis_fallback(default_return=None) - def add_email_register_error_rate_limit(email: str) -> None: - key = f"email_register_error_rate_limit:{email}" - count = redis_client.get(key) - if count is None: - count = 0 - count = int(count) + 1 - redis_client.setex(key, dify_config.EMAIL_REGISTER_LOCKOUT_DURATION, count) - @staticmethod @redis_fallback(default_return=False) def is_forgot_password_error_rate_limit(email: str) -> bool: @@ -1095,24 +1027,6 @@ class AccountService: key = f"forgot_password_error_rate_limit:{email}" redis_client.delete(key) - @staticmethod - @redis_fallback(default_return=False) - def is_email_register_error_rate_limit(email: str) -> bool: - key = f"email_register_error_rate_limit:{email}" - count = redis_client.get(key) - if count is None: - return False - count = int(count) - if count > AccountService.EMAIL_REGISTER_MAX_ERROR_LIMITS: - return True - return False - - @staticmethod - @redis_fallback(default_return=None) - def reset_email_register_error_rate_limit(email: str): - key = f"email_register_error_rate_limit:{email}" - redis_client.delete(key) - @staticmethod @redis_fallback(default_return=None) def add_change_email_error_rate_limit(email: str): diff --git a/api/services/entities/account_entities.py b/api/services/entities/account_entities.py index 21cfc2df00d..b53a739eba2 100644 --- a/api/services/entities/account_entities.py +++ b/api/services/entities/account_entities.py @@ -116,6 +116,30 @@ class AccountEmailResetResult: account: AccountSnapshot | None = None +class AccountEmailRegistrationPhase(StrEnum): + REGISTER = "register" + + +@dataclass(frozen=True, slots=True) +class AccountEmailRegistrationToken: + email: str + code: str + phase: AccountEmailRegistrationPhase | None = None + + +@dataclass(frozen=True, slots=True) +class AccountEmailRegistrationVerification: + email: str + token: str + + +@dataclass(frozen=True, slots=True) +class AccountSessionTokens: + access_token: str + refresh_token: str + csrf_token: str + + class AccountChangeEmailPhase(StrEnum): OLD_EMAIL = "old_email" OLD_EMAIL_VERIFIED = "old_email_verified" 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 7eee2e68102..7b5f859877b 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 @@ -1,30 +1,57 @@ -"""Unit tests for email register controller endpoints.""" +"""Unit tests for the email-registration Flask adapter.""" from __future__ import annotations -from collections.abc import Callable -from unittest.mock import MagicMock, patch +from collections.abc import Callable, Generator +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock, patch import pytest from flask import Flask +from pydantic import ValidationError +from controllers.console import bp as console_bp from controllers.console.auth.email_register import ( EmailRegisterCheckApi, EmailRegisterResetApi, + EmailRegisterResetPayload, EmailRegisterSendEmailApi, ) -from controllers.console.auth.error import NormalizedEmailAlreadyInUseError -from controllers.console.error import AccountInFreezeError, EmailDomainSuspendedError +from controllers.console.auth.error import ( + EmailAlreadyInUseError, + EmailCodeError, + EmailRegisterLimitError, + EmailRegisterRateLimitExceededError, + InvalidEmailError, + InvalidTokenError, + NormalizedEmailAlreadyInUseError, + PasswordMismatchError, +) +from controllers.console.error import ( + AccountInFreezeError, + EmailDomainSuspendedError, + EmailSendIpLimitError, + SeatsLimitExceeded, +) from enums import DeploymentEdition -from models.account import Account -from services.entities.feature_entities import SystemFeatureModel -from services.errors.account import ( +from services.account_email_registration_service import AccountEmailRegistrationService +from services.account_errors import ( + AccountEmailAlreadyInUseError, + AccountEmailDomainSuspendedError, + AccountEmailFrozenError, AccountNormalizedEmailAlreadyInUseError, - AccountRegisterError, -) -from services.errors.account import ( - EmailDomainSuspendedError as EmailDomainSuspendedRegistrationError, + EmailRegistrationPasswordMismatchError, + EmailRegistrationSeatsLimitError, + EmailRegistrationSendIPLimitedError, + EmailRegistrationSendRateLimitError, + EmailRegistrationVerificationLimitError, + InvalidEmailRegistrationAddressError, + InvalidEmailRegistrationCodeError, + InvalidEmailRegistrationTokenError, ) +from services.entities.account_entities import AccountEmailRegistrationVerification, AccountSessionTokens +from services.entities.feature_entities import SystemFeatureModel @pytest.fixture(autouse=True) @@ -32,6 +59,33 @@ def _cloud_edition(config_overrides: Callable[..., None]) -> None: config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) +@contextmanager +def _request( + app: Flask, + service: Mock, + *, + path: str, + payload: dict[str, str], +) -> Generator[None, None, None]: + services = SimpleNamespace(accounts=SimpleNamespace(email_registration=service)) + features = SystemFeatureModel( + deployment_edition=DeploymentEdition.CLOUD, + enable_email_password_login=True, + is_allow_register=True, + ) + with ( + patch("controllers.console.auth.email_register.application_services", return_value=services), + patch("controllers.console.flask_admission.FeatureService.get_system_features", return_value=features), + patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1"), + app.test_request_context(path, method="POST", json=payload), + ): + yield + + +def _service() -> Mock: + return Mock(spec=AccountEmailRegistrationService) + + def test_normalized_email_conflict_exposes_a_distinct_error_code() -> None: error = NormalizedEmailAlreadyInUseError() @@ -40,321 +94,210 @@ def test_normalized_email_conflict_exposes_a_distinct_error_code() -> None: assert error.data["code"] == "normalized_email_already_in_use" -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.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( - self, - mock_extract_ip, - mock_is_email_send_ip_limit, - mock_is_freeze, - mock_send_mail, - mock_get_account, - app: Flask, +def test_send_email_delegates_with_remote_ip(app: Flask) -> None: + service = _service() + service.send_code.return_value = "token-123" + + with _request( + app, + service, + path="/email-register/send-email", + payload={"email": "Invitee@Example.com", "language": "zh-Hans"}, ): - mock_send_mail.return_value = "token-123" - mock_is_freeze.return_value = False - account = Account(name="Invitee", email="invitee@example.com") - mock_get_account.return_value = account + response = EmailRegisterSendEmailApi().post() - feature_flags = SystemFeatureModel( - deployment_edition=DeploymentEdition.COMMUNITY, - enable_email_password_login=True, - is_allow_register=True, - ) - with ( - 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", "language": "en-US"}, - ): - response = EmailRegisterSendEmailApi().post() + assert response == {"result": "success", "data": "token-123"} + assert service.send_code.call_args.kwargs == { + "remote_ip": "127.0.0.1", + "requested_email": "Invitee@Example.com", + "requested_language": "zh-Hans", + } - assert response == {"result": "success", "data": "token-123"} - mock_is_freeze.assert_called_once_with("invitee@example.com") - mock_send_mail.assert_called_once_with(email="invitee@example.com", account=account, language="en-US") - 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), - ], +@pytest.mark.parametrize( + ("service_error", "http_error"), + [ + pytest.param(EmailRegistrationSendIPLimitedError(), EmailSendIpLimitError, id="ip-limit"), + pytest.param(EmailRegistrationSendRateLimitError(1), EmailRegisterRateLimitExceededError, id="send-limit"), + pytest.param(AccountEmailFrozenError(), AccountInFreezeError, id="frozen"), + pytest.param(AccountEmailDomainSuspendedError(), EmailDomainSuspendedError, id="suspended-domain"), + ], +) +def test_send_email_translates_application_errors( + app: Flask, + service_error: Exception, + http_error: type[Exception], +) -> None: + service = _service() + service.send_code.side_effect = service_error + + with _request( + app, + service, + path="/email-register/send-email", + payload={"email": "invitee@example.com"}, + ): + with pytest.raises(http_error): + EmailRegisterSendEmailApi().post() + + +def test_verify_email_code_serializes_application_result(app: Flask) -> None: + service = _service() + service.verify_code.return_value = AccountEmailRegistrationVerification( + email="user@example.com", + token="verified-token", ) - @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, + + with _request( + app, + service, + path="/email-register/validity", + payload={"email": "User@Example.com", "code": "123456", "token": "pending-token"}, ): - mock_get_freeze_type.return_value = freeze_type - feature_flags = SystemFeatureModel( - deployment_edition=DeploymentEdition.COMMUNITY, - enable_email_password_login=True, - is_allow_register=True, - ) + response = EmailRegisterCheckApi().post() - with ( - 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") - @patch("controllers.console.auth.email_register.AccountService.generate_email_register_token") - @patch("controllers.console.auth.email_register.AccountService.revoke_email_register_token") - @patch("controllers.console.auth.email_register.AccountService.add_email_register_error_rate_limit") - @patch("controllers.console.auth.email_register.AccountService.get_email_register_data") - @patch("controllers.console.auth.email_register.AccountService.is_email_register_error_rate_limit") - def test_validity_normalizes_email_before_checks( - self, - mock_rate_limit_check, - mock_get_data, - mock_add_rate, - mock_revoke, - mock_generate_token, - mock_reset_rate, - app: Flask, - ): - mock_rate_limit_check.return_value = False - mock_get_data.return_value = {"email": "User@Example.com", "code": "4321"} - mock_generate_token.return_value = (None, "new-token") - - feature_flags = SystemFeatureModel( - deployment_edition=DeploymentEdition.COMMUNITY, - enable_email_password_login=True, - is_allow_register=True, - ) - with ( - patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags), - ): - with app.test_request_context( - "/email-register/validity", - method="POST", - json={"email": "User@Example.com", "code": "4321", "token": "token-123"}, - ): - response = EmailRegisterCheckApi().post() - - assert response == {"is_valid": True, "email": "user@example.com", "token": "new-token"} - mock_rate_limit_check.assert_called_once_with("user@example.com") - mock_generate_token.assert_called_once_with( - "user@example.com", code="4321", additional_data={"phase": "register"} - ) - mock_reset_rate.assert_called_once_with("user@example.com") - mock_add_rate.assert_not_called() - mock_revoke.assert_called_once_with("token-123") - - -class TestEmailRegisterResetApi: - @pytest.mark.parametrize( - ("service_error", "expected_error"), - [ - (EmailDomainSuspendedRegistrationError(), EmailDomainSuspendedError), - (AccountNormalizedEmailAlreadyInUseError(), NormalizedEmailAlreadyInUseError), - (AccountRegisterError("frozen"), AccountInFreezeError), - ], + assert response == {"is_valid": True, "email": "user@example.com", "token": "verified-token"} + service.verify_code.assert_called_once_with( + email="User@Example.com", + code="123456", + token="pending-token", ) - @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, + + +@pytest.mark.parametrize( + ("service_error", "http_error"), + [ + pytest.param(EmailRegistrationVerificationLimitError(), EmailRegisterLimitError, id="attempt-limit"), + pytest.param(InvalidEmailRegistrationTokenError(), InvalidTokenError, id="token"), + pytest.param(InvalidEmailRegistrationAddressError(), InvalidEmailError, id="email"), + pytest.param(InvalidEmailRegistrationCodeError(), EmailCodeError, id="code"), + ], +) +def test_verify_email_code_translates_application_errors( + app: Flask, + service_error: Exception, + http_error: type[Exception], +) -> None: + service = _service() + service.verify_code.side_effect = service_error + + with _request( + app, + service, + path="/email-register/validity", + payload={"email": "user@example.com", "code": "wrong", "token": "pending-token"}, ): - mock_create_account.side_effect = service_error + with pytest.raises(http_error): + EmailRegisterCheckApi().post() - 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") - @patch("controllers.console.auth.email_register.AccountService.get_account_by_email_with_case_fallback") - @patch("controllers.console.auth.email_register.AccountService.revoke_email_register_token") - @patch("controllers.console.auth.email_register.AccountService.get_email_register_data") - @patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1") - def test_reset_creates_account_with_normalized_email( - self, - mock_extract_ip, - mock_get_data, - mock_revoke_token, - mock_get_account, - mock_create_account, - mock_login, - mock_reset_login_rate, - app: Flask, +def test_register_delegates_and_serializes_tokens(app: Flask) -> None: + service = _service() + service.register.return_value = AccountSessionTokens( + access_token="access", + refresh_token="refresh", + csrf_token="csrf", + ) + + with _request( + app, + service, + path="/email-register", + payload={ + "token": "verified-token", + "new_password": "ValidPass123!", + "password_confirm": "ValidPass123!", + "language": "zh-Hans", + "timezone": "Asia/Shanghai", + }, ): - mock_get_data.return_value = {"phase": "register", "email": "Invitee@Example.com"} - mock_create_account.return_value = Account(name="Invitee", email="invitee@example.com") - token_pair = MagicMock() - token_pair.model_dump.return_value = {"access_token": "a", "refresh_token": "r"} - mock_login.return_value = token_pair - mock_get_account.return_value = None + response = EmailRegisterResetApi().post() - feature_flags = SystemFeatureModel( - deployment_edition=DeploymentEdition.COMMUNITY, - enable_email_password_login=True, - is_allow_register=True, - ) - with ( - patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags), - ): - with app.test_request_context( - "/email-register", - method="POST", - json={"token": "token-123", "new_password": "ValidPass123!", "password_confirm": "ValidPass123!"}, - ): - response = EmailRegisterResetApi().post() + assert response == { + "result": "success", + "data": {"access_token": "access", "refresh_token": "refresh", "csrf_token": "csrf"}, + } + assert service.register.call_args.kwargs == { + "remote_ip": "127.0.0.1", + "token": "verified-token", + "new_password": "ValidPass123!", + "password_confirm": "ValidPass123!", + "language": "zh-Hans", + "timezone": "Asia/Shanghai", + } - assert response == {"result": "success", "data": {"access_token": "a", "refresh_token": "r"}} - mock_create_account.assert_called_once_with( - email="invitee@example.com", - password="ValidPass123!", - timezone=None, - language=None, - ip_address="127.0.0.1", - ) - mock_reset_login_rate.assert_called_once_with("invitee@example.com") - mock_revoke_token.assert_called_once_with("token-123") - mock_extract_ip.assert_called_once() - @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") - @patch("controllers.console.auth.email_register.AccountService.get_account_by_email_with_case_fallback") - @patch("controllers.console.auth.email_register.AccountService.revoke_email_register_token") - @patch("controllers.console.auth.email_register.AccountService.get_email_register_data") - @patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1") - def test_reset_passes_timezone_to_new_account( - self, - mock_extract_ip, - mock_get_data, - mock_revoke_token, - mock_get_account, - mock_create_account, - mock_login, - mock_reset_login_rate, - app: Flask, +@pytest.mark.parametrize( + ("service_error", "http_error"), + [ + pytest.param(EmailRegistrationPasswordMismatchError(), PasswordMismatchError, id="password"), + pytest.param(InvalidEmailRegistrationTokenError(), InvalidTokenError, id="token"), + pytest.param( + AccountNormalizedEmailAlreadyInUseError(), + NormalizedEmailAlreadyInUseError, + id="normalized-email-in-use", + ), + pytest.param(AccountEmailAlreadyInUseError(), EmailAlreadyInUseError, id="email-in-use"), + pytest.param(EmailRegistrationSeatsLimitError(), SeatsLimitExceeded, id="seat-limit"), + pytest.param(AccountEmailFrozenError(), AccountInFreezeError, id="frozen"), + pytest.param(AccountEmailDomainSuspendedError(), EmailDomainSuspendedError, id="suspended-domain"), + ], +) +def test_register_translates_application_errors( + app: Flask, + service_error: Exception, + http_error: type[Exception], +) -> None: + service = _service() + service.register.side_effect = service_error + + with _request( + app, + service, + path="/email-register", + payload={ + "token": "verified-token", + "new_password": "ValidPass123!", + "password_confirm": "ValidPass123!", + }, ): - mock_get_data.return_value = {"phase": "register", "email": "Invitee@Example.com"} - mock_create_account.return_value = Account(name="Invitee", email="invitee@example.com") - token_pair = MagicMock() - token_pair.model_dump.return_value = {"access_token": "a", "refresh_token": "r"} - mock_login.return_value = token_pair - mock_get_account.return_value = None + with pytest.raises(http_error): + EmailRegisterResetApi().post() - feature_flags = SystemFeatureModel( - deployment_edition=DeploymentEdition.COMMUNITY, - enable_email_password_login=True, - is_allow_register=True, + +def test_reset_payload_rejects_invalid_timezone() -> None: + with pytest.raises(ValidationError): + EmailRegisterResetPayload.model_validate( + { + "token": "token-123", + "new_password": "ValidPass123!", + "password_confirm": "ValidPass123!", + "timezone": "", + } ) - with ( - patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags), - ): - with app.test_request_context( - "/email-register", - method="POST", - json={ - "token": "token-123", - "new_password": "ValidPass123!", - "password_confirm": "ValidPass123!", - "timezone": "Asia/Shanghai", - }, - ): - response = EmailRegisterResetApi().post() - assert response == {"result": "success", "data": {"access_token": "a", "refresh_token": "r"}} - mock_create_account.assert_called_once_with( - email="invitee@example.com", - password="ValidPass123!", - timezone="Asia/Shanghai", - language=None, - ip_address="127.0.0.1", + +def test_invalid_password_is_sanitized_by_real_error_handler(caplog: pytest.LogCaptureFixture) -> None: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(console_bp) + features = SystemFeatureModel( + deployment_edition=DeploymentEdition.CLOUD, + enable_email_password_login=True, + is_allow_register=True, + ) + password_marker = "SecretMarker" + + with patch("controllers.console.flask_admission.FeatureService.get_system_features", return_value=features): + response = app.test_client().post( + "/console/api/email-register", + json={ + "token": "verified-token", + "new_password": password_marker, + "password_confirm": password_marker, + }, ) - mock_reset_login_rate.assert_called_once_with("invitee@example.com") - mock_revoke_token.assert_called_once_with("token-123") - mock_extract_ip.assert_called_once() - @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") - @patch("controllers.console.auth.email_register.AccountService.get_account_by_email_with_case_fallback") - @patch("controllers.console.auth.email_register.AccountService.revoke_email_register_token") - @patch("controllers.console.auth.email_register.AccountService.get_email_register_data") - @patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1") - def test_reset_passes_language_to_new_account( - self, - mock_extract_ip, - mock_get_data, - mock_revoke_token, - mock_get_account, - mock_create_account, - mock_login, - mock_reset_login_rate, - app: Flask, - ): - mock_get_data.return_value = {"phase": "register", "email": "Invitee@Example.com"} - mock_create_account.return_value = Account(name="Invitee", email="invitee@example.com") - token_pair = MagicMock() - token_pair.model_dump.return_value = {"access_token": "a", "refresh_token": "r"} - mock_login.return_value = token_pair - mock_get_account.return_value = None - - feature_flags = SystemFeatureModel( - deployment_edition=DeploymentEdition.COMMUNITY, - enable_email_password_login=True, - is_allow_register=True, - ) - with ( - patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags), - ): - with app.test_request_context( - "/email-register", - method="POST", - json={ - "token": "token-123", - "new_password": "ValidPass123!", - "password_confirm": "ValidPass123!", - "language": "zh-Hans", - }, - ): - response = EmailRegisterResetApi().post() - - assert response == {"result": "success", "data": {"access_token": "a", "refresh_token": "r"}} - mock_create_account.assert_called_once_with( - email="invitee@example.com", - password="ValidPass123!", - timezone=None, - language="zh-Hans", - ip_address="127.0.0.1", - ) - mock_reset_login_rate.assert_called_once_with("invitee@example.com") - mock_revoke_token.assert_called_once_with("token-123") - mock_extract_ip.assert_called_once() + assert response.status_code == 422 + assert password_marker not in response.get_data(as_text=True) + assert password_marker not in caplog.text diff --git a/api/tests/unit_tests/controllers/console/auth/test_email_register_language.py b/api/tests/unit_tests/controllers/console/auth/test_email_register_language.py deleted file mode 100644 index e8331bda8cc..00000000000 --- a/api/tests/unit_tests/controllers/console/auth/test_email_register_language.py +++ /dev/null @@ -1,44 +0,0 @@ -from unittest.mock import ANY, patch - -import pytest -from pydantic import ValidationError - -from controllers.console.auth.email_register import EmailRegisterResetApi, EmailRegisterResetPayload -from models.account import Account - - -@patch("controllers.console.auth.email_register.AccountService.create_account_and_tenant") -def test_create_new_account_uses_requested_language(mock_create_account): - account = Account(name="Invitee", email="invitee@example.com") - mock_create_account.return_value = account - - result = EmailRegisterResetApi()._create_new_account( - "invitee@example.com", - "ValidPass123!", - timezone="Asia/Shanghai", - language="zh-Hans", - ) - - assert result is account - mock_create_account.assert_called_once_with( - email="invitee@example.com", - name="invitee@example.com", - password="ValidPass123!", - interface_language="zh-Hans", - timezone="Asia/Shanghai", - ip_address=None, - check_normalized_email=True, - session=ANY, - ) - - -def test_reset_payload_rejects_invalid_timezone(): - with pytest.raises(ValidationError): - EmailRegisterResetPayload.model_validate( - { - "token": "token-123", - "new_password": "ValidPass123!", - "password_confirm": "ValidPass123!", - "timezone": "", - } - ) diff --git a/api/tests/unit_tests/controllers/console/test_wraps.py b/api/tests/unit_tests/controllers/console/test_wraps.py index be562c20133..20a7291ec66 100644 --- a/api/tests/unit_tests/controllers/console/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/test_wraps.py @@ -196,6 +196,64 @@ class TestCurrentContextInjection: login_required.assert_called_once() account_initialization_required.assert_called_once() + def test_console_email_registration_admission_checks_features_once(self): + features = SimpleNamespace(enable_email_password_login=True, is_allow_register=True) + with ( + patch( + "controllers.console.flask_admission.setup_required", side_effect=lambda view: view + ) as setup_required, + patch( + "controllers.console.flask_admission.FeatureService.get_system_features", + return_value=features, + ) as get_system_features, + ): + + class Handler: + @flask_admission.console_email_registration_admission + def post(self): + return "ok" + + with Flask(__name__).test_request_context(): + result = Handler().post() + + assert result == "ok" + setup_required.assert_called_once() + get_system_features.assert_called_once_with() + + @pytest.mark.parametrize( + ("enable_email_password_login", "is_allow_register"), + [ + pytest.param(False, True, id="password-login-disabled"), + pytest.param(True, False, id="registration-disabled"), + ], + ) + def test_console_email_registration_admission_rejects_disabled_features( + self, + enable_email_password_login: bool, + is_allow_register: bool, + ) -> None: + features = SimpleNamespace( + enable_email_password_login=enable_email_password_login, + is_allow_register=is_allow_register, + ) + with ( + patch("controllers.console.flask_admission.setup_required", side_effect=lambda view: view), + patch( + "controllers.console.flask_admission.FeatureService.get_system_features", + return_value=features, + ), + ): + + class Handler: + @flask_admission.console_email_registration_admission + def post(self): + return "ok" + + with Flask(__name__).test_request_context(), pytest.raises(HTTPException) as exc_info: + Handler().post() + + assert exc_info.value.code == 403 + def test_console_account_admission_preserves_route_kwarg_named_request_context(self): current_user = make_account() 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 1cf77a1f27b..ee3d2125204 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -32,6 +32,12 @@ from services.account_activation_adapters import ( RegisterServiceInvitationTokenStore, ) from services.account_avatar_file_gateway import SQLAlchemyAccountAvatarFileGateway +from services.account_email_registration_adapters import ( + AccountServiceRegistrationGateway, + BillingAccountRegistrationPolicyGateway, + RedisEmailRegistrationSecurityGateway, + TokenManagerEmailRegistrationTokenGateway, +) from services.app_site_service import AppSiteService from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService from services.billing_portal_service import BillingPortalService @@ -367,6 +373,13 @@ def test_build_application_services_wires_account_profile_repository( assert services.accounts.initialization._accounts is accounts assert not services.accounts.initialization._invitation_required assert services.accounts.change_email._accounts is accounts + email_registration = services.accounts.email_registration + assert email_registration._accounts is accounts + assert isinstance(email_registration._tokens, TokenManagerEmailRegistrationTokenGateway) + assert isinstance(email_registration._security, RedisEmailRegistrationSecurityGateway) + assert isinstance(email_registration._account_policy, BillingAccountRegistrationPolicyGateway) + assert isinstance(email_registration._registration, AccountServiceRegistrationGateway) + assert email_registration._registration._session_factory is sqlite_session_factory assert services.accounts.education._accounts is accounts assert services.accounts.deletion._accounts is accounts assert services.accounts.deletion._memberships is services.workspace_queries._workspaces diff --git a/api/tests/unit_tests/repositories/test_account_repository.py b/api/tests/unit_tests/repositories/test_account_repository.py index daab5e6e7c3..945d2c7ddd6 100644 --- a/api/tests/unit_tests/repositories/test_account_repository.py +++ b/api/tests/unit_tests/repositories/test_account_repository.py @@ -97,6 +97,20 @@ def test_account_repository_updates_password( assert persisted.password_salt == "new-salt" +def test_account_repository_finds_email_with_lowercase_fallback( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + _persist_account(sqlite_session) + repository = SQLAlchemyAccountRepository(sqlite_session_factory) + + account = repository.find_by_email("Account@Example.com") + + assert account is not None + assert account.id == "account-1" + assert account.email == "account@example.com" + + def test_account_integration_repository_lists_integrations( sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], diff --git a/api/tests/unit_tests/services/test_account_email_registration_adapters.py b/api/tests/unit_tests/services/test_account_email_registration_adapters.py new file mode 100644 index 00000000000..14d2fd725f2 --- /dev/null +++ b/api/tests/unit_tests/services/test_account_email_registration_adapters.py @@ -0,0 +1,176 @@ +from unittest.mock import Mock, patch + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from extensions.ext_redis import RedisClientWrapper +from models.account import Account +from services.account_email_registration_adapters import ( + AccountServiceRegistrationGateway, + BillingAccountRegistrationPolicyGateway, + RedisEmailRegistrationSecurityGateway, + TokenManagerEmailRegistrationTokenGateway, +) +from services.account_errors import ( + AccountEmailDomainSuspendedError, + AccountNormalizedEmailAlreadyInUseError, + EmailRegistrationSeatsLimitError, +) +from services.account_service import TokenPair +from services.entities.account_entities import AccountEmailRegistrationPhase, AccountEmailRegistrationToken +from services.errors.account import ( + AccountNormalizedEmailAlreadyInUseError as AccountNormalizedEmailAlreadyInUseServiceError, +) +from services.errors.account import EmailDomainSuspendedError, SeatsLimitExceededError + + +def test_token_gateway_rejects_malformed_payload() -> None: + gateway = TokenManagerEmailRegistrationTokenGateway() + + with patch( + "services.account_email_registration_adapters.TokenManager.get_token_data", + return_value={"email": "user@example.com", "phase": "unknown"}, + ): + assert gateway.get("token") is None + + +def test_token_gateway_issues_verified_registration_state() -> None: + gateway = TokenManagerEmailRegistrationTokenGateway() + token_data = AccountEmailRegistrationToken( + email="user@example.com", + code="123456", + phase=AccountEmailRegistrationPhase.REGISTER, + ) + + with patch( + "services.account_email_registration_adapters.TokenManager.generate_token", + return_value="token", + ) as generate_token: + assert gateway.issue(token_data) == "token" + + generate_token.assert_called_once_with( + email="user@example.com", + token_type="email_register", + additional_data={"code": "123456", "phase": "register"}, + ) + + +def test_security_gateway_delegates_ip_limit_to_existing_policy_owner() -> None: + redis = Mock(spec=RedisClientWrapper) + gateway = RedisEmailRegistrationSecurityGateway( + redis=redis, + verification_failure_limit=5, + verification_lockout_duration=600, + ) + + with patch( + "services.account_email_registration_adapters.AccountService.is_email_send_ip_limit", + return_value=False, + ) as is_email_send_ip_limit: + assert gateway.is_ip_limited("127.0.0.1") is False + + is_email_send_ip_limit.assert_called_once_with("127.0.0.1") + redis.get.assert_not_called() + + +def test_security_gateway_uses_registration_and_login_keys() -> None: + redis = Mock(spec=RedisClientWrapper) + redis.get.return_value = 1 + gateway = RedisEmailRegistrationSecurityGateway( + redis=redis, + verification_failure_limit=5, + verification_lockout_duration=600, + ) + + with patch( + "services.account_email_registration_adapters.AccountService.reset_login_error_rate_limit" + ) as reset_login_error_rate_limit: + gateway.record_verification_failure("user@example.com") + gateway.reset_verification_failures("user@example.com") + gateway.reset_login_failures("user@example.com") + + redis.setex.assert_called_once_with("email_register_error_rate_limit:user@example.com", 600, 2) + redis.delete.assert_called_once_with("email_register_error_rate_limit:user@example.com") + reset_login_error_rate_limit.assert_called_once_with("user@example.com") + + +def test_billing_policy_is_disabled_outside_cloud() -> None: + gateway = BillingAccountRegistrationPolicyGateway(enabled=False) + + with patch("services.account_email_registration_adapters.BillingService.get_email_freeze_type") as freeze_type: + assert gateway.get_freeze_type("user@example.com") is None + + freeze_type.assert_not_called() + + +@pytest.mark.parametrize( + ("service_error", "application_error"), + [ + pytest.param(SeatsLimitExceededError(), EmailRegistrationSeatsLimitError, id="seat-limit"), + pytest.param(EmailDomainSuspendedError(), AccountEmailDomainSuspendedError, id="suspended-domain"), + pytest.param( + AccountNormalizedEmailAlreadyInUseServiceError(), + AccountNormalizedEmailAlreadyInUseError, + id="normalized-email-in-use", + ), + ], +) +def test_registration_gateway_translates_account_provisioning_errors( + sqlite_session_factory: sessionmaker[Session], + service_error: Exception, + application_error: type[Exception], +) -> None: + gateway = AccountServiceRegistrationGateway(session_factory=sqlite_session_factory) + + with patch( + "services.account_email_registration_adapters.AccountService.create_account_and_tenant", + side_effect=service_error, + ): + with pytest.raises(application_error): + gateway.create( + email="user@example.com", + password="ValidPass123!", + interface_language="en-US", + timezone=None, + ip_address="127.0.0.1", + ) + + +def test_registration_gateway_owns_short_lived_sessions( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + gateway = AccountServiceRegistrationGateway(session_factory=sqlite_session_factory) + + def create_account(*, session: Session, **_: object) -> Account: + account = Account(name="user@example.com", email="user@example.com") + account.id = "account-1" + session.add(account) + session.commit() + return account + + with patch( + "services.account_email_registration_adapters.AccountService.create_account_and_tenant", + side_effect=create_account, + ) as create_account_and_tenant: + account_id = gateway.create( + email="user@example.com", + password="ValidPass123!", + interface_language="en-US", + timezone=None, + ip_address="127.0.0.1", + ) + + assert create_account_and_tenant.call_args.kwargs["check_normalized_email"] is True + sqlite_session.expire_all() + assert sqlite_session.get(Account, account_id) is not None + + with patch( + "services.account_email_registration_adapters.AccountService.login", + return_value=TokenPair(access_token="access", refresh_token="refresh", csrf_token="csrf"), + ) as login: + tokens = gateway.login(account_id, ip_address="127.0.0.1") + + assert tokens.access_token == "access" + assert login.call_args.kwargs["account"].id == account_id + assert isinstance(login.call_args.kwargs["session"], Session) diff --git a/api/tests/unit_tests/services/test_account_email_registration_service.py b/api/tests/unit_tests/services/test_account_email_registration_service.py new file mode 100644 index 00000000000..b3091f51438 --- /dev/null +++ b/api/tests/unit_tests/services/test_account_email_registration_service.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +from datetime import datetime +from unittest.mock import Mock + +import pytest + +from services.account_email_registration_service import ( + AccountEmailRegistrationService, + AccountRegistrationGateway, + AccountRegistrationPolicyGateway, + EmailRegistrationCodeGenerator, + EmailRegistrationNotificationGateway, + EmailRegistrationSecurityGateway, + EmailRegistrationSendLimiter, + EmailRegistrationTokenGateway, +) +from services.account_errors import ( + AccountEmailAlreadyInUseError, + AccountEmailDomainSuspendedError, + EmailRegistrationPasswordMismatchError, + InvalidEmailRegistrationCodeError, + InvalidEmailRegistrationTokenError, +) +from services.account_ports import AccountRepository +from services.entities.account_entities import ( + AccountEmailRegistrationPhase, + AccountEmailRegistrationToken, + AccountSessionTokens, + AccountSnapshot, +) + + +def _account(*, email: str = "stored@example.com") -> AccountSnapshot: + return AccountSnapshot( + id="account-1", + name="Stored Account", + email=email, + avatar=None, + is_password_set=True, + interface_language="en-US", + interface_theme="light", + timezone="UTC", + last_login_at=None, + last_login_ip=None, + status="active", + initialized_at=datetime(2026, 1, 1), + created_at=datetime(2026, 1, 1), + ) + + +def _service() -> tuple[AccountEmailRegistrationService, dict[str, Mock]]: + dependencies = { + "accounts": Mock(spec=AccountRepository), + "tokens": Mock(spec=EmailRegistrationTokenGateway), + "codes": Mock(spec=EmailRegistrationCodeGenerator), + "notifications": Mock(spec=EmailRegistrationNotificationGateway), + "send_limits": Mock(spec=EmailRegistrationSendLimiter), + "security": Mock(spec=EmailRegistrationSecurityGateway), + "account_policy": Mock(spec=AccountRegistrationPolicyGateway), + "registration": Mock(spec=AccountRegistrationGateway), + } + service = AccountEmailRegistrationService( + accounts=dependencies["accounts"], + tokens=dependencies["tokens"], + codes=dependencies["codes"], + notifications=dependencies["notifications"], + send_limits=dependencies["send_limits"], + security=dependencies["security"], + account_policy=dependencies["account_policy"], + registration=dependencies["registration"], + ) + dependencies["accounts"].find_by_email.return_value = None + dependencies["codes"].generate.return_value = "123456" + dependencies["tokens"].issue.return_value = "token-1" + dependencies["send_limits"].is_limited.return_value = False + dependencies["security"].is_ip_limited.return_value = False + dependencies["security"].is_verification_limited.return_value = False + dependencies["account_policy"].get_freeze_type.return_value = None + return service, dependencies + + +def test_send_code_uses_case_fallback_account_and_existing_account_notification() -> None: + service, dependencies = _service() + dependencies["accounts"].find_by_email.return_value = _account(email="Stored@Example.com") + + token = service.send_code( + remote_ip="127.0.0.1", + requested_email="Stored@Example.com", + requested_language="zh-Hans", + ) + + assert token == "token-1" + dependencies["accounts"].find_by_email.assert_called_once_with("Stored@Example.com") + dependencies["tokens"].issue.assert_called_once_with( + AccountEmailRegistrationToken(email="Stored@Example.com", code="123456") + ) + dependencies["notifications"].send_account_exists.assert_called_once_with( + email="Stored@Example.com", + account_name="Stored Account", + language="zh-Hans", + ) + dependencies["send_limits"].record.assert_called_once_with("Stored@Example.com") + + +def test_send_code_normalizes_new_account_email_and_language() -> None: + service, dependencies = _service() + + service.send_code( + remote_ip="127.0.0.1", + requested_email="New@Example.com", + requested_language="unsupported", + ) + + dependencies["notifications"].send_code.assert_called_once_with( + email="new@example.com", + code="123456", + language="en-US", + ) + + +def test_send_code_rejects_suspended_domain_before_account_lookup() -> None: + service, dependencies = _service() + dependencies["account_policy"].get_freeze_type.return_value = "email_domain_suspended" + + with pytest.raises(AccountEmailDomainSuspendedError): + service.send_code( + remote_ip="127.0.0.1", + requested_email="user@suspended.example", + requested_language=None, + ) + + dependencies["accounts"].find_by_email.assert_not_called() + + +def test_verify_code_rotates_token_into_register_phase() -> None: + service, dependencies = _service() + dependencies["tokens"].get.return_value = AccountEmailRegistrationToken( + email="User@Example.com", + code="123456", + ) + dependencies["tokens"].issue.return_value = "verified-token" + + verification = service.verify_code( + email="USER@example.com", + code="123456", + token="pending-token", + ) + + assert verification.email == "user@example.com" + assert verification.token == "verified-token" + dependencies["tokens"].revoke.assert_called_once_with("pending-token") + dependencies["tokens"].issue.assert_called_once_with( + AccountEmailRegistrationToken( + email="user@example.com", + code="123456", + phase=AccountEmailRegistrationPhase.REGISTER, + ) + ) + dependencies["security"].reset_verification_failures.assert_called_once_with("user@example.com") + + +def test_verify_code_records_failure_without_consuming_token() -> None: + service, dependencies = _service() + dependencies["tokens"].get.return_value = AccountEmailRegistrationToken( + email="user@example.com", + code="123456", + ) + + with pytest.raises(InvalidEmailRegistrationCodeError): + service.verify_code(email="user@example.com", code="wrong", token="pending-token") + + dependencies["security"].record_verification_failure.assert_called_once_with("user@example.com") + dependencies["tokens"].revoke.assert_not_called() + + +def test_register_creates_account_and_logs_it_in() -> None: + service, dependencies = _service() + dependencies["tokens"].get.return_value = AccountEmailRegistrationToken( + email="New@Example.com", + code="123456", + phase=AccountEmailRegistrationPhase.REGISTER, + ) + dependencies["registration"].create.return_value = "account-1" + expected_tokens = AccountSessionTokens(access_token="access", refresh_token="refresh", csrf_token="csrf") + dependencies["registration"].login.return_value = expected_tokens + + tokens = service.register( + remote_ip="127.0.0.1", + token="verified-token", + new_password="ValidPass123!", + password_confirm="ValidPass123!", + language="zh-Hans", + timezone="Asia/Shanghai", + ) + + assert tokens == expected_tokens + dependencies["tokens"].revoke.assert_called_once_with("verified-token") + dependencies["accounts"].find_by_email.assert_called_once_with("New@Example.com") + dependencies["registration"].create.assert_called_once_with( + email="new@example.com", + password="ValidPass123!", + interface_language="zh-Hans", + timezone="Asia/Shanghai", + ip_address="127.0.0.1", + ) + dependencies["registration"].login.assert_called_once_with("account-1", ip_address="127.0.0.1") + dependencies["security"].reset_login_failures.assert_called_once_with("new@example.com") + + +def test_register_rejects_password_mismatch_before_reading_token() -> None: + service, dependencies = _service() + + with pytest.raises(EmailRegistrationPasswordMismatchError): + service.register( + remote_ip="127.0.0.1", + token="verified-token", + new_password="ValidPass123!", + password_confirm="DifferentPass123!", + language=None, + timezone=None, + ) + + dependencies["tokens"].get.assert_not_called() + + +def test_register_requires_verified_registration_phase() -> None: + service, dependencies = _service() + dependencies["tokens"].get.return_value = AccountEmailRegistrationToken( + email="new@example.com", + code="123456", + ) + + with pytest.raises(InvalidEmailRegistrationTokenError): + service.register( + remote_ip="127.0.0.1", + token="pending-token", + new_password="ValidPass123!", + password_confirm="ValidPass123!", + language=None, + timezone=None, + ) + + dependencies["tokens"].revoke.assert_not_called() + + +def test_register_consumes_token_before_rejecting_existing_account() -> None: + service, dependencies = _service() + dependencies["tokens"].get.return_value = AccountEmailRegistrationToken( + email="existing@example.com", + code="123456", + phase=AccountEmailRegistrationPhase.REGISTER, + ) + dependencies["accounts"].find_by_email.return_value = _account(email="existing@example.com") + + with pytest.raises(AccountEmailAlreadyInUseError): + service.register( + remote_ip="127.0.0.1", + token="verified-token", + new_password="ValidPass123!", + password_confirm="ValidPass123!", + language=None, + timezone=None, + ) + + dependencies["tokens"].revoke.assert_called_once_with("verified-token") + dependencies["registration"].create.assert_not_called() From 7167d7d66289b8e748a5fa89528956abf700a8ce Mon Sep 17 00:00:00 2001 From: Wu Tianwei <30284043+WTW0313@users.noreply.github.com> Date: Mon, 31 Aug 2026 03:33:09 +0000 Subject: [PATCH 06/21] fix(rbac): grant agent access to all legacy roles (#41511) --- api/services/enterprise/rbac_service.py | 2 ++ .../unit_tests/services/enterprise/test_rbac_service.py | 8 ++------ 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/api/services/enterprise/rbac_service.py b/api/services/enterprise/rbac_service.py index b3d1bfac0fd..bf95b8650b1 100644 --- a/api/services/enterprise/rbac_service.py +++ b/api/services/enterprise/rbac_service.py @@ -464,6 +464,7 @@ _LEGACY_WORKSPACE_NORMAL_KEYS: list[str] = [ "plugin.install", "credential.use", "app_library.access", + "agent.manage", ] _LEGACY_WORKSPACE_DATASET_OPERATOR_KEYS: list[str] = [ @@ -471,6 +472,7 @@ _LEGACY_WORKSPACE_DATASET_OPERATOR_KEYS: list[str] = [ "plugin.install", "dataset.create_and_management", "dataset.external.connect", + "agent.manage", ] _LEGACY_APP_OWNER_KEYS: list[str] = [ diff --git a/api/tests/unit_tests/services/enterprise/test_rbac_service.py b/api/tests/unit_tests/services/enterprise/test_rbac_service.py index 432466a9975..0ce4390fd0c 100644 --- a/api/tests/unit_tests/services/enterprise/test_rbac_service.py +++ b/api/tests/unit_tests/services/enterprise/test_rbac_service.py @@ -1128,16 +1128,12 @@ class TestListOption: class TestLegacyAgentManageKey: def test_legacy_agent_manage_key_membership(self): - # Mirrors the builtin roles in the rbac service, which grant agent.manage - # to owner/admin/editor only. + # Preserve Agent access for every legacy role while external RBAC is disabled. for keys in ( svc._LEGACY_WORKSPACE_OWNER_KEYS, svc._LEGACY_WORKSPACE_ADMIN_KEYS, svc._LEGACY_WORKSPACE_EDITOR_KEYS, - ): - assert "agent.manage" in keys - for keys in ( svc._LEGACY_WORKSPACE_NORMAL_KEYS, svc._LEGACY_WORKSPACE_DATASET_OPERATOR_KEYS, ): - assert "agent.manage" not in keys + assert "agent.manage" in keys From c6ca39504219fce2600e2d28d0974685bf1f71a0 Mon Sep 17 00:00:00 2001 From: Coding On Star <447357187@qq.com> Date: Mon, 31 Aug 2026 03:39:37 +0000 Subject: [PATCH 07/21] fix(web): make registration tracking consent-safe and retryable (#41365) Co-authored-by: CodingOnStar --- .../__tests__/external-service-sync.spec.tsx | 106 ++++ .../(commonLayout)/external-service-sync.tsx | 28 +- .../oauth-registration-analytics.spec.tsx | 169 +++++- .../__tests__/registration-tracking.spec.ts | 574 ++++++++++++++---- .../registration-consent-coordinator.tsx | 15 + .../amplitude/registration-session-state.ts | 99 +++ .../base/amplitude/registration-tracking.ts | 303 +++++++-- web/app/components/base/amplitude/utils.ts | 7 +- .../__tests__/analytics-disabled.spec.tsx | 26 + .../__tests__/analytics-runtimes.spec.tsx | 6 + .../__tests__/cloud-analytics.spec.tsx | 25 +- .../analytics-consent/analytics-disabled.tsx | 14 + .../analytics-consent/cloud-analytics.tsx | 3 +- .../base/analytics-consent/consent-store.ts | 2 +- .../console-analytics-runtime.tsx | 2 + .../oauth-registration-analytics.tsx | 92 ++- .../__tests__/console-bootstrap.spec.tsx | 2 + web/service/__tests__/base-request.spec.ts | 54 ++ web/service/base.ts | 12 + web/service/common.spec.ts | 35 ++ web/service/use-common.ts | 2 + 21 files changed, 1336 insertions(+), 240 deletions(-) create mode 100644 web/app/(commonLayout)/__tests__/external-service-sync.spec.tsx create mode 100644 web/app/components/base/amplitude/registration-consent-coordinator.tsx create mode 100644 web/app/components/base/amplitude/registration-session-state.ts create mode 100644 web/app/components/base/analytics-consent/__tests__/analytics-disabled.spec.tsx create mode 100644 web/app/components/base/analytics-consent/analytics-disabled.tsx diff --git a/web/app/(commonLayout)/__tests__/external-service-sync.spec.tsx b/web/app/(commonLayout)/__tests__/external-service-sync.spec.tsx new file mode 100644 index 00000000000..c851acaf03b --- /dev/null +++ b/web/app/(commonLayout)/__tests__/external-service-sync.spec.tsx @@ -0,0 +1,106 @@ +import { render, waitFor } from '@testing-library/react' +import { REGISTRATION_SUCCESS_STORAGE_KEY } from '@/app/components/base/amplitude/registration-session-state' +import { rememberRegistrationSuccess } from '@/app/components/base/amplitude/registration-tracking' +import { AmplitudeIdentitySync } from '../external-service-sync' + +const { mockSetUserId, mockSetUserProperties, mockTrackEvent } = vi.hoisted(() => ({ + mockSetUserId: vi.fn(), + mockSetUserProperties: vi.fn(), + mockTrackEvent: vi.fn((..._args: unknown[]) => ({ + promise: Promise.resolve({ code: 200 }), + })), +})) + +vi.mock('@tanstack/react-query', async (importOriginal) => { + const original = await importOriginal() + return { + ...original, + useSuspenseQuery: () => ({ + data: { + id: 'account-id', + email: 'person@example.com', + name: 'Person', + is_password_set: true, + }, + }), + } +}) + +vi.mock('jotai', async (importOriginal) => { + const original = await importOriginal() + return { + ...original, + useAtomValue: () => ({ + id: 'workspace-id', + name: 'Workspace', + plan: 'professional', + role: 'owner', + }), + } +}) + +vi.mock('@/features/account-profile/client', () => ({ + userProfileQueryOptions: () => ({}), +})) + +vi.mock('@/app/components/base/amplitude', () => ({ + setUserId: (...args: unknown[]) => mockSetUserId(...args), + setUserProperties: (...args: unknown[]) => mockSetUserProperties(...args), +})) + +vi.mock('@/app/components/base/amplitude/utils', () => ({ + trackEvent: (...args: unknown[]) => mockTrackEvent(...args), +})) + +vi.mock('@/app/components/base/amplitude/init', () => ({ + getIsAmplitudeInitialized: () => true, +})) + +vi.mock('@/app/components/base/analytics-consent/consent-store', async (importOriginal) => { + const original = + await importOriginal() + return { + ...original, + getAnalyticsConsent: () => 'granted', + } +}) + +describe('AmplitudeIdentitySync', () => { + beforeEach(() => { + vi.clearAllMocks() + window.sessionStorage.clear() + vi.spyOn(globalThis.crypto, 'randomUUID').mockReturnValue( + '11111111-1111-4111-8111-111111111111', + ) + }) + + it('sets identity before flushing a marker that already exists', async () => { + rememberRegistrationSuccess({ method: 'oauth' }) + + render() + + await waitFor(() => expect(mockTrackEvent).toHaveBeenCalledTimes(1)) + expect(mockSetUserId).toHaveBeenCalledWith('person@example.com') + expect(mockSetUserProperties).toHaveBeenCalledTimes(1) + expect(mockSetUserId.mock.invocationCallOrder[0]).toBeLessThan( + mockTrackEvent.mock.invocationCallOrder[0]!, + ) + expect(mockSetUserProperties.mock.invocationCallOrder[0]).toBeLessThan( + mockTrackEvent.mock.invocationCallOrder[0]!, + ) + }) + + it('flushes a marker created after identity sync without repeating unchanged identity updates', async () => { + render() + + await waitFor(() => expect(mockSetUserId).toHaveBeenCalledTimes(1)) + expect(mockTrackEvent).not.toHaveBeenCalled() + + rememberRegistrationSuccess({ method: 'email' }) + + await waitFor(() => expect(mockTrackEvent).toHaveBeenCalledTimes(1)) + expect(mockSetUserId).toHaveBeenCalledTimes(1) + expect(mockSetUserProperties).toHaveBeenCalledTimes(1) + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + }) +}) diff --git a/web/app/(commonLayout)/external-service-sync.tsx b/web/app/(commonLayout)/external-service-sync.tsx index cec7155626e..08bb786d379 100644 --- a/web/app/(commonLayout)/external-service-sync.tsx +++ b/web/app/(commonLayout)/external-service-sync.tsx @@ -5,9 +5,13 @@ import type { DeploymentEdition } from '@dify/contracts/api/console/system-featu import type { GetWorkspacesCurrentSummaryResponse } from '@dify/contracts/api/console/workspaces/types.gen' import { skipToken, useQuery, useSuspenseQuery } from '@tanstack/react-query' import { useAtomValue } from 'jotai' -import { Fragment, useEffect, useRef } from 'react' +import { Fragment, useEffect, useRef, useSyncExternalStore } from 'react' import { setUserId, setUserProperties } from '@/app/components/base/amplitude' -import { flushRegistrationSuccess } from '@/app/components/base/amplitude/registration-tracking' +import { + flushRegistrationSuccess, + getRegistrationSuccessSnapshot, + subscribeRegistrationSuccess, +} from '@/app/components/base/amplitude/registration-tracking' import { useAmplitudeInitialized } from '@/app/components/base/amplitude/use-amplitude-initialized' import { useAnalyticsConsent } from '@/app/components/base/analytics-consent/consent-store' import { setZendeskConversationFields } from '@/app/components/base/zendesk/utils' @@ -43,13 +47,18 @@ function buildAmplitudeProperties({ return properties } -function AmplitudeIdentitySync() { +export function AmplitudeIdentitySync() { const { data: userProfile } = useSuspenseQuery({ ...userProfileQueryOptions(), select: (data) => data.profile, }) const currentWorkspace = useAtomValue(currentWorkspaceAtom) const lastIdentityRef = useRef(undefined) + const registrationSnapshot = useSyncExternalStore( + subscribeRegistrationSuccess, + getRegistrationSuccessSnapshot, + getRegistrationSuccessSnapshot, + ) useEffect(() => { if (!userProfile.id) return @@ -63,13 +72,14 @@ function AmplitudeIdentitySync() { properties, }) - if (identity === lastIdentityRef.current) return + if (identity !== lastIdentityRef.current) { + setUserId(userProfile.email) + setUserProperties(properties) + lastIdentityRef.current = identity + } - setUserId(userProfile.email) - setUserProperties(properties) - flushRegistrationSuccess() - lastIdentityRef.current = identity - }, [currentWorkspace, userProfile]) + void flushRegistrationSuccess() + }, [currentWorkspace, registrationSnapshot, userProfile]) return null } diff --git a/web/app/components/__tests__/oauth-registration-analytics.spec.tsx b/web/app/components/__tests__/oauth-registration-analytics.spec.tsx index eb8f7deacdf..53d18f0d2b9 100644 --- a/web/app/components/__tests__/oauth-registration-analytics.spec.tsx +++ b/web/app/components/__tests__/oauth-registration-analytics.spec.tsx @@ -1,12 +1,31 @@ import { render, waitFor } from '@testing-library/react' import Cookies from 'js-cookie' +import { StrictMode } from 'react' import { beforeEach, describe, expect, it, vi } from 'vite-plus/test' import { useSearchParams } from '@/next/navigation' import { OAuthRegistrationAnalytics } from '../oauth-registration-analytics' -const { mockSendGAEvent, mockRememberRegistrationSuccess } = vi.hoisted(() => ({ - mockSendGAEvent: vi.fn(), +const { + mockConsent, + mockNormalizeRegistrationAttribution, + mockRememberRegistrationSuccess, + mockSendGAEvent, +} = vi.hoisted(() => ({ + mockConsent: { value: 'granted' as 'unknown' | 'denied' | 'granted' | 'disabled' }, + mockNormalizeRegistrationAttribution: vi.fn((value: Record | null) => { + if (!value) return null + const allowed = Object.fromEntries( + Object.entries(value).filter( + ([key, item]) => + ['utm_source', 'utm_medium', 'utm_campaign', 'utm_content', 'utm_term', 'slug'].includes( + key, + ) && typeof item === 'string', + ), + ) + return Object.keys(allowed).length ? allowed : null + }), mockRememberRegistrationSuccess: vi.fn(), + mockSendGAEvent: vi.fn(), })) vi.mock('@/utils/gtag', () => ({ @@ -17,7 +36,14 @@ vi.mock('@/next/navigation', () => ({ useSearchParams: vi.fn(), })) +vi.mock('../base/analytics-consent/consent-store', () => ({ + useAnalyticsConsent: () => mockConsent.value, +})) + vi.mock('../base/amplitude/registration-tracking', () => ({ + normalizeRegistrationAttribution: ( + ...args: Parameters + ) => mockNormalizeRegistrationAttribution(...args), rememberRegistrationSuccess: (...args: unknown[]) => mockRememberRegistrationSuccess(...args), })) @@ -33,22 +59,74 @@ const setSearchParams = (searchParams = '') => { describe('OAuthRegistrationAnalytics', () => { beforeEach(() => { vi.clearAllMocks() + window.sessionStorage.clear() + mockConsent.value = 'granted' + mockRememberRegistrationSuccess.mockReturnValue(true) Cookies.remove('utm_info') vi.spyOn(console, 'error').mockImplementation(() => {}) setSearchParams() }) - it('should track oauth registration with utm info and clear the query flag', async () => { + it('queues the Amplitude marker while consent is unknown and cleans the URL after persist', async () => { + mockConsent.value = 'unknown' + Cookies.set('utm_info', JSON.stringify({ utm_source: 'linkedin', slug: 'agent-launch' })) + setSearchParams('oauth_new_user=true&source=signin') + + render() + + await waitFor(() => { + expect(mockRememberRegistrationSuccess).toHaveBeenCalledWith({ + method: 'oauth', + utmInfo: { utm_source: 'linkedin', slug: 'agent-launch' }, + }) + }) + expect(mockSendGAEvent).toHaveBeenCalledTimes(1) + expect(Cookies.get('utm_info')).toBeUndefined() + expect(window.location.search).toBe('?source=signin') + }) + + it('keeps the recoverable OAuth signal when marker persistence fails', () => { + Cookies.set('utm_info', JSON.stringify({ utm_source: 'linkedin' })) + setSearchParams('oauth_new_user=true&source=signin') + mockRememberRegistrationSuccess.mockReturnValue(false) + + render() + + expect(mockRememberRegistrationSuccess).toHaveBeenCalledTimes(1) + expect(window.location.search).toBe('?oauth_new_user=true&source=signin') + expect(Cookies.get('utm_info')).toBeTruthy() + }) + + it('keeps the OAuth marker while consent is unknown, then cleans without a second Amplitude queue on denial', async () => { + mockConsent.value = 'unknown' + Cookies.set('utm_info', JSON.stringify({ utm_source: 'blog' })) + setSearchParams('oauth_new_user=true') + + const { rerender } = render() + + await waitFor(() => expect(mockRememberRegistrationSuccess).toHaveBeenCalledTimes(1)) + expect(mockSendGAEvent).toHaveBeenCalledTimes(1) + + mockConsent.value = 'denied' + rerender() + + await waitFor(() => expect(window.location.search).toBe('')) + expect(mockRememberRegistrationSuccess).toHaveBeenCalledTimes(1) + expect(mockSendGAEvent).toHaveBeenCalledTimes(1) + expect(Cookies.get('utm_info')).toBeUndefined() + }) + + it('queues immediately with pre-granted consent and keeps only allowlisted UTM fields', async () => { Cookies.set( 'utm_info', JSON.stringify({ utm_source: 'linkedin', slug: 'agent-launch', + arbitrary: 'discard-me', + utm_term: { nested: true }, }), ) - setSearchParams('oauth_new_user=true&source=signin') - const replaceStateSpy = vi.spyOn(window.history, 'replaceState') render() @@ -64,16 +142,13 @@ describe('OAuthRegistrationAnalytics', () => { slug: 'agent-launch', }) expect(Cookies.get('utm_info')).toBeUndefined() - - await waitFor(() => { - expect(replaceStateSpy).toHaveBeenCalledWith(null, '', '/signin?source=signin') - }) + expect(window.location.search).toBe('?source=signin') }) - it('should fall back to the base registration event when the utm cookie is invalid', async () => { + it('uses the base event and cleans up when the UTM cookie is malformed', async () => { Cookies.set('utm_info', '{invalid-json') - setSearchParams('oauth_new_user=true') + render() await waitFor(() => { @@ -89,23 +164,77 @@ describe('OAuthRegistrationAnalytics', () => { expect(Cookies.get('utm_info')).toBeUndefined() }) - it('should do nothing without the oauth registration query flag', () => { + it('cleans a false OAuth marker immediately without tracking or clearing utm_info', async () => { + mockConsent.value = 'unknown' + Cookies.set('utm_info', JSON.stringify({ utm_source: 'blog' })) + setSearchParams('oauth_new_user=false') + + render() + + await waitFor(() => expect(window.location.search).toBe('')) + expect(mockRememberRegistrationSuccess).not.toHaveBeenCalled() + expect(mockSendGAEvent).not.toHaveBeenCalled() + expect(Cookies.get('utm_info')).toBe(JSON.stringify({ utm_source: 'blog' })) + }) + + it('tracks GA and Amplitude once across StrictMode effects and rerenders', async () => { + setSearchParams('oauth_new_user=true') + + const { rerender } = render( + + + , + ) + + rerender( + + + , + ) + + await waitFor(() => expect(mockRememberRegistrationSuccess).toHaveBeenCalledTimes(1)) + expect(mockSendGAEvent).toHaveBeenCalledTimes(1) + }) + + it('tracks GA once across an unknown-consent remount that simulates reload', async () => { + mockConsent.value = 'unknown' + setSearchParams('oauth_new_user=true') + + const firstRender = render() + await waitFor(() => expect(mockRememberRegistrationSuccess).toHaveBeenCalledTimes(1)) + firstRender.unmount() + render() + + expect(mockSendGAEvent).toHaveBeenCalledTimes(1) + }) + + it('treats analytics-disabled consent as terminal and cleans without Amplitude', async () => { + mockConsent.value = 'disabled' + Cookies.set('utm_info', JSON.stringify({ utm_source: 'blog' })) + setSearchParams('oauth_new_user=true') + + render() + + await waitFor(() => expect(window.location.search).toBe('')) + expect(mockRememberRegistrationSuccess).not.toHaveBeenCalled() + expect(Cookies.get('utm_info')).toBeUndefined() + }) + + it('does nothing without the OAuth registration query marker', () => { render() expect(mockRememberRegistrationSuccess).not.toHaveBeenCalled() expect(mockSendGAEvent).not.toHaveBeenCalled() }) - it('should clear a false oauth registration query flag without tracking', async () => { - setSearchParams('oauth_new_user=false') - const replaceStateSpy = vi.spyOn(window.history, 'replaceState') + it('clears an abandoned flow guard so a later OAuth registration can emit GA', () => { + window.sessionStorage.setItem('oauth_registration_ga_sent', 'true') + const abandonedFlow = render() + abandonedFlow.unmount() + setSearchParams('oauth_new_user=true') render() - await waitFor(() => { - expect(replaceStateSpy).toHaveBeenCalledWith(null, '', '/signin') - }) - expect(mockRememberRegistrationSuccess).not.toHaveBeenCalled() - expect(mockSendGAEvent).not.toHaveBeenCalled() + expect(mockSendGAEvent).toHaveBeenCalledTimes(1) }) }) diff --git a/web/app/components/base/amplitude/__tests__/registration-tracking.spec.ts b/web/app/components/base/amplitude/__tests__/registration-tracking.spec.ts index 30c6707a702..e5979a79623 100644 --- a/web/app/components/base/amplitude/__tests__/registration-tracking.spec.ts +++ b/web/app/components/base/amplitude/__tests__/registration-tracking.spec.ts @@ -1,55 +1,182 @@ import { - flushRegistrationSuccess, + discardRegistrationSessionState, REGISTRATION_SUCCESS_STORAGE_KEY, +} from '../registration-session-state' +import { + coordinateRegistrationConsent, + flushRegistrationSuccess, rememberRegistrationSuccess, + subscribeRegistrationSuccess, } from '../registration-tracking' const mockTrackEvent = vi.hoisted(() => vi.fn()) +const mockAmplitudeInitialized = vi.hoisted(() => ({ value: true })) const mockConsent = vi.hoisted(() => ({ - value: 'granted' as 'unknown' | 'denied' | 'granted', + value: 'granted' as 'unknown' | 'denied' | 'granted' | 'disabled', })) vi.mock('../utils', () => ({ trackEvent: (...args: unknown[]) => mockTrackEvent(...args), })) +vi.mock('../init', () => ({ + getIsAmplitudeInitialized: () => mockAmplitudeInitialized.value, +})) + vi.mock('@/app/components/base/analytics-consent/consent-store', () => ({ getAnalyticsConsent: () => mockConsent.value, })) +const successResult = () => ({ + promise: Promise.resolve({ code: 200 }), +}) + +const getStoredMarker = () => + JSON.parse(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)!) + describe('registration tracking', () => { beforeEach(() => { vi.clearAllMocks() vi.unstubAllGlobals() + vi.useRealTimers() window.sessionStorage.clear() mockConsent.value = 'granted' + mockAmplitudeInitialized.value = true + mockTrackEvent.mockImplementation(successResult) + coordinateRegistrationConsent('denied') + mockConsent.value = 'granted' + vi.spyOn(globalThis.crypto, 'randomUUID').mockReturnValue( + '11111111-1111-4111-8111-111111111111', + ) }) - // Captures the registration event for a later flush instead of firing it right away. - describe('rememberRegistrationSuccess', () => { - it('should store the base event and not track immediately when there is no utm info', () => { - rememberRegistrationSuccess({ method: 'email' }) + afterEach(() => { + discardRegistrationSessionState() + vi.useRealTimers() + }) - expect(JSON.parse(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)!)).toEqual({ - eventName: 'user_registration_success', - properties: { method: 'email' }, + describe('rememberRegistrationSuccess', () => { + it('stores a versioned marker with stable delivery metadata and allowlisted attribution', () => { + vi.setSystemTime(new Date('2026-08-26T09:00:00.000Z')) + + const persisted = rememberRegistrationSuccess({ + method: 'email', + utmInfo: { + utm_source: 'linkedin', + utm_medium: 'social', + utm_campaign: 'launch', + utm_content: 'hero', + utm_term: 'agents', + slug: 'agent-launch', + unexpected: 'discard-me', + nested: { unsafe: true }, + }, + }) + + const occurredAt = Date.now() + expect(getStoredMarker()).toEqual({ + version: 2, + registrationId: '11111111-1111-4111-8111-111111111111', + occurredAt, + expiresAt: occurredAt + 24 * 60 * 60 * 1000, + eventName: 'user_registration_success_with_utm', + method: 'email', + attribution: { + utm_source: 'linkedin', + utm_medium: 'social', + utm_campaign: 'launch', + utm_content: 'hero', + utm_term: 'agents', + slug: 'agent-launch', + }, }) expect(mockTrackEvent).not.toHaveBeenCalled() + expect(persisted).toBe(true) }) - it('should store the utm event and merge utm info into properties when utm info is present', () => { - rememberRegistrationSuccess({ - method: 'oauth', - utmInfo: { utm_source: 'linkedin', slug: 'agent-launch' }, + it('persists the latest email marker while consent is unknown so a later grant can flush it', async () => { + mockConsent.value = 'unknown' + rememberRegistrationSuccess({ method: 'email', utmInfo: { utm_source: 'first' } }) + rememberRegistrationSuccess({ method: 'email', utmInfo: { utm_source: 'latest' } }) + + expect(getStoredMarker()).toMatchObject({ + version: 2, + method: 'email', + attribution: { utm_source: 'latest' }, }) - expect(JSON.parse(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)!)).toEqual({ - eventName: 'user_registration_success_with_utm', - properties: { method: 'oauth', utm_source: 'linkedin', slug: 'agent-launch' }, - }) + await flushRegistrationSuccess() + expect(mockTrackEvent).not.toHaveBeenCalled() + + mockConsent.value = 'granted' + await flushRegistrationSuccess() + + expect(mockTrackEvent).toHaveBeenCalledTimes(1) + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() }) - it('should swallow errors when writing to sessionStorage fails', () => { + it('discards an unknown-consent marker on denial before a later grant can flush it', async () => { + mockConsent.value = 'unknown' + rememberRegistrationSuccess({ method: 'email' }) + + coordinateRegistrationConsent('denied') + mockConsent.value = 'granted' + await flushRegistrationSuccess() + + expect(mockTrackEvent).not.toHaveBeenCalled() + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + }) + + it('discards a pending marker and GA guard at an account boundary', async () => { + mockConsent.value = 'unknown' + rememberRegistrationSuccess({ method: 'email' }) + window.sessionStorage.setItem('oauth_registration_ga_sent', 'true') + + discardRegistrationSessionState() + mockConsent.value = 'granted' + await flushRegistrationSuccess() + + expect(mockTrackEvent).not.toHaveBeenCalled() + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + expect(window.sessionStorage.getItem('oauth_registration_ga_sent')).toBeNull() + }) + + it.each(['denied', 'disabled'] as const)( + 'discards a stored marker when consent changes to %s', + (consent) => { + rememberRegistrationSuccess({ method: 'email' }) + + coordinateRegistrationConsent(consent) + + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + }, + ) + + it('persists an oauth marker while consent is unknown so a reload can still flush it', async () => { + mockConsent.value = 'unknown' + rememberRegistrationSuccess({ method: 'oauth' }) + + expect(getStoredMarker()).toMatchObject({ method: 'oauth' }) + + mockConsent.value = 'granted' + await flushRegistrationSuccess() + + expect(mockTrackEvent).toHaveBeenCalledTimes(1) + }) + + it('notifies consumers only after a marker is stored', () => { + const listener = vi.fn() + const unsubscribe = subscribeRegistrationSuccess(listener) + + rememberRegistrationSuccess({ method: 'email' }) + + expect(listener).toHaveBeenCalledTimes(1) + unsubscribe() + }) + + it('swallows sessionStorage write errors without notifying consumers', () => { + const listener = vi.fn() + const unsubscribe = subscribeRegistrationSuccess(listener) vi.stubGlobal('window', { sessionStorage: { getItem: vi.fn(() => null), @@ -60,158 +187,365 @@ describe('registration tracking', () => { }, }) - try { - expect(() => rememberRegistrationSuccess({ method: 'email' })).not.toThrow() - } finally { - vi.unstubAllGlobals() - } + expect(rememberRegistrationSuccess({ method: 'email' })).toBe(false) + expect(listener).not.toHaveBeenCalled() + unsubscribe() + }) + }) + + describe('flushRegistrationSuccess', () => { + it('waits for a successful SDK result before acknowledging the marker', async () => { + let resolveTrack!: (result: { code: number }) => void + mockTrackEvent.mockReturnValue({ + promise: new Promise((resolve) => { + resolveTrack = resolve + }), + }) + vi.setSystemTime(new Date('2026-08-26T09:00:00.000Z')) + rememberRegistrationSuccess({ method: 'oauth', utmInfo: { utm_source: 'blog' } }) + + const flushPromise = flushRegistrationSuccess() + + expect(getStoredMarker()).toBeTruthy() + expect(mockTrackEvent).toHaveBeenCalledWith( + 'user_registration_success_with_utm', + { + method: 'oauth', + utm_source: 'blog', + registration_id: '11111111-1111-4111-8111-111111111111', + event_version: 2, + tracking_contract_version: 'consent_wait_v2', + }, + { + insert_id: '11111111-1111-4111-8111-111111111111', + time: Date.now(), + }, + ) + + resolveTrack({ code: 200 }) + await flushPromise + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() }) - it.each(['unknown', 'denied'] as const)( - 'should not cache an event while consent is %s', - (consent) => { - mockConsent.value = consent + it.each([ + ['unknown consent', () => (mockConsent.value = 'unknown')], + ['uninitialized Amplitude', () => (mockAmplitudeInitialized.value = false)], + ])('defers without deleting for %s', async (_label, makeIneligible) => { + rememberRegistrationSuccess({ method: 'email' }) + makeIneligible() + + await flushRegistrationSuccess() + + expect(mockTrackEvent).not.toHaveBeenCalled() + expect(getStoredMarker()).toBeTruthy() + }) + + it('discards a pending marker when consent is denied', async () => { + rememberRegistrationSuccess({ method: 'oauth' }) + mockConsent.value = 'denied' + + await flushRegistrationSuccess() + + expect(mockTrackEvent).not.toHaveBeenCalled() + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + }) + + it('retains the marker on SDK rejection or a non-success result and reuses its id and time', async () => { + vi.setSystemTime(new Date('2026-08-26T09:00:00.000Z')) + rememberRegistrationSuccess({ method: 'email' }) + const marker = getStoredMarker() + mockTrackEvent + .mockReturnValueOnce({ promise: Promise.reject(new Error('network failed')) }) + .mockReturnValueOnce({ promise: Promise.resolve({ code: 500 }) }) + .mockImplementation(successResult) + + await flushRegistrationSuccess() + await flushRegistrationSuccess() + + expect(getStoredMarker()).toEqual(marker) + expect(mockTrackEvent.mock.calls[0]?.[2]).toEqual({ + insert_id: marker.registrationId, + time: marker.occurredAt, + }) + expect(mockTrackEvent.mock.calls[1]?.[2]).toEqual({ + insert_id: marker.registrationId, + time: marker.occurredAt, + }) + }) + + it.each(['rejected acknowledgement', 'non-success acknowledgement'] as const)( + 'continues with a replacement marker after a %s', + async (oldAcknowledgement) => { + const firstRegistrationId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa' + const replacementRegistrationId = 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb' + vi.spyOn(globalThis.crypto, 'randomUUID') + .mockReturnValueOnce(firstRegistrationId) + .mockReturnValueOnce(replacementRegistrationId) + let resolveOldAcknowledgement!: (result: { code: number }) => void + let rejectOldAcknowledgement!: (error: Error) => void + const pendingOldAcknowledgement = new Promise<{ code: number }>((resolve, reject) => { + resolveOldAcknowledgement = resolve + rejectOldAcknowledgement = reject + }) + mockTrackEvent + .mockReturnValueOnce({ promise: pendingOldAcknowledgement }) + .mockImplementation(successResult) rememberRegistrationSuccess({ method: 'email' }) + const firstFlush = flushRegistrationSuccess() + rememberRegistrationSuccess({ method: 'email' }) + const replacementFlush = flushRegistrationSuccess() + expect(replacementFlush).toBe(firstFlush) + expect(mockTrackEvent).toHaveBeenCalledTimes(1) + + if (oldAcknowledgement === 'rejected acknowledgement') + rejectOldAcknowledgement(new Error('network failed')) + else resolveOldAcknowledgement({ code: 500 }) + await firstFlush + + expect(mockTrackEvent).toHaveBeenCalledTimes(2) + expect(mockTrackEvent.mock.calls[1]?.[2]).toEqual({ + insert_id: replacementRegistrationId, + time: expect.any(Number), + }) expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() }, ) - }) - // Replays the remembered event exactly once, after the user ID has been attached. - describe('flushRegistrationSuccess', () => { - it('should track the remembered event and clear it from storage', () => { - rememberRegistrationSuccess({ method: 'email', utmInfo: { utm_source: 'blog' } }) - - flushRegistrationSuccess() + it('retries an acknowledgement-failed marker after the backoff delay', async () => { + vi.useFakeTimers() + rememberRegistrationSuccess({ method: 'email' }) + const marker = getStoredMarker() + mockTrackEvent + .mockReturnValueOnce({ promise: Promise.reject(new Error('network failed')) }) + .mockImplementation(successResult) + await flushRegistrationSuccess() + expect(getStoredMarker()).toEqual(marker) expect(mockTrackEvent).toHaveBeenCalledTimes(1) - expect(mockTrackEvent).toHaveBeenCalledWith('user_registration_success_with_utm', { - method: 'email', - utm_source: 'blog', + + await vi.advanceTimersByTimeAsync(1000) + + expect(mockTrackEvent).toHaveBeenCalledTimes(2) + expect(mockTrackEvent.mock.calls[1]?.[2]).toEqual({ + insert_id: marker.registrationId, + time: marker.occurredAt, }) expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() }) - it('should do nothing when there is no pending event', () => { - flushRegistrationSuccess() + it('does not retry an acknowledgement-failed marker after an account boundary', async () => { + rememberRegistrationSuccess({ method: 'email' }) + mockTrackEvent.mockReturnValueOnce({ promise: Promise.reject(new Error('network failed')) }) + await flushRegistrationSuccess() + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).not.toBeNull() + + discardRegistrationSessionState() + await flushRegistrationSuccess() + + expect(mockTrackEvent).toHaveBeenCalledTimes(1) + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + }) + + it.each([ + ['session discard', 'resolves'] as const, + ['session discard', 'rejects'] as const, + ['denied consent', 'resolves'] as const, + ['disabled analytics', 'resolves'] as const, + ])( + 'isolates a new registration flush after %s while the old SDK acknowledgement %s', + async (invalidation, oldAcknowledgement) => { + const accountARegistrationId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa' + const accountBRegistrationId = 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb' + vi.spyOn(globalThis.crypto, 'randomUUID') + .mockReturnValueOnce(accountARegistrationId) + .mockReturnValueOnce(accountBRegistrationId) + + let resolveAccountA!: (result: { code: number }) => void + let rejectAccountA!: (error: Error) => void + let resolveAccountB!: (result: { code: number }) => void + const accountAAcknowledgement = new Promise<{ code: number }>((resolve, reject) => { + resolveAccountA = resolve + rejectAccountA = reject + }) + const accountBAcknowledgement = new Promise<{ code: number }>((resolve) => { + resolveAccountB = resolve + }) + mockTrackEvent + .mockReturnValueOnce({ promise: accountAAcknowledgement }) + .mockReturnValueOnce({ promise: accountBAcknowledgement }) + + rememberRegistrationSuccess({ method: 'email' }) + const accountAFlush = flushRegistrationSuccess() + + if (invalidation === 'session discard') { + discardRegistrationSessionState() + } else { + const terminalConsent = invalidation === 'denied consent' ? 'denied' : 'disabled' + mockConsent.value = terminalConsent + coordinateRegistrationConsent(terminalConsent) + mockConsent.value = 'granted' + } + + rememberRegistrationSuccess({ method: 'email' }) + const accountBMarker = window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY) + const accountBFlush = flushRegistrationSuccess() + + const settleAccountA = () => { + if (oldAcknowledgement === 'resolves') resolveAccountA({ code: 200 }) + else rejectAccountA(new Error('account A request failed')) + } + + try { + expect(accountBMarker).not.toBeNull() + expect(accountBFlush).not.toBe(accountAFlush) + expect(mockTrackEvent).toHaveBeenCalledTimes(2) + expect(mockTrackEvent.mock.calls.map((call) => call[2])).toEqual([ + { insert_id: accountARegistrationId, time: expect.any(Number) }, + { insert_id: accountBRegistrationId, time: expect.any(Number) }, + ]) + + settleAccountA() + await accountAFlush + + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBe( + accountBMarker, + ) + expect(flushRegistrationSuccess()).toBe(accountBFlush) + expect(mockTrackEvent).toHaveBeenCalledTimes(2) + + resolveAccountB({ code: 200 }) + await accountBFlush + + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + } finally { + settleAccountA() + resolveAccountB({ code: 200 }) + await Promise.allSettled([accountAFlush, accountBFlush]) + } + }, + ) + + it('coalesces concurrent flushes into one SDK send', async () => { + let resolveTrack!: (result: { code: number }) => void + mockTrackEvent.mockReturnValue({ + promise: new Promise((resolve) => { + resolveTrack = resolve + }), + }) + rememberRegistrationSuccess({ method: 'email' }) + + const first = flushRegistrationSuccess() + const second = flushRegistrationSuccess() + + expect(mockTrackEvent).toHaveBeenCalledTimes(1) + resolveTrack({ code: 200 }) + await Promise.all([first, second]) + }) + + it('discards expired and malformed markers without tracking', async () => { + vi.setSystemTime(new Date('2026-08-26T09:00:00.000Z')) + rememberRegistrationSuccess({ method: 'email' }) + vi.setSystemTime(new Date('2026-08-27T09:00:00.000Z')) + + await flushRegistrationSuccess() + expect(mockTrackEvent).not.toHaveBeenCalled() + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + + const malformedMarkers = [ + '{not-json', + JSON.stringify({ version: 1, eventName: 'user_registration_success' }), + JSON.stringify({ + version: 2, + registrationId: 'id', + occurredAt: Date.now(), + expiresAt: Date.now() + 1000, + eventName: 'arbitrary_event', + method: 'email', + attribution: {}, + }), + ] + + for (const raw of malformedMarkers) { + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, raw) + await flushRegistrationSuccess() + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + } expect(mockTrackEvent).not.toHaveBeenCalled() }) - it('should fire the event at most once across repeated flushes', () => { - rememberRegistrationSuccess({ method: 'oauth' }) + it('accepts a persisted timestamp just inside the five-minute clock-skew allowance', async () => { + vi.setSystemTime(new Date('2026-08-26T09:04:59.999Z')) + rememberRegistrationSuccess({ method: 'email' }) + vi.setSystemTime(new Date('2026-08-26T09:00:00.000Z')) - flushRegistrationSuccess() - flushRegistrationSuccess() + await flushRegistrationSuccess() expect(mockTrackEvent).toHaveBeenCalledTimes(1) }) - it('should discard a pending event when consent was revoked before flush', () => { - rememberRegistrationSuccess({ method: 'oauth' }) - mockConsent.value = 'denied' + it('rejects a persisted timestamp just outside the five-minute clock-skew allowance', async () => { + vi.setSystemTime(new Date('2026-08-26T09:05:00.001Z')) + rememberRegistrationSuccess({ method: 'email' }) + vi.setSystemTime(new Date('2026-08-26T09:00:00.000Z')) - flushRegistrationSuccess() + await flushRegistrationSuccess() expect(mockTrackEvent).not.toHaveBeenCalled() expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() }) - it('should clear malformed pending data without tracking', () => { - window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, '{not-json') - - flushRegistrationSuccess() - - expect(mockTrackEvent).not.toHaveBeenCalled() - expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() - }) - - it('should clear the pending entry without tracking when it has no event name', () => { - window.sessionStorage.setItem( - REGISTRATION_SUCCESS_STORAGE_KEY, - JSON.stringify({ properties: { method: 'email' } }), - ) - - flushRegistrationSuccess() - - expect(mockTrackEvent).not.toHaveBeenCalled() - expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() - }) - - it('should stop without tracking when reading from sessionStorage throws', () => { + it('handles storage read errors without throwing', async () => { vi.stubGlobal('window', { sessionStorage: { getItem: () => { throw new Error('read failed') }, setItem: vi.fn(), - removeItem: vi.fn(), - }, - }) - - try { - expect(() => flushRegistrationSuccess()).not.toThrow() - expect(mockTrackEvent).not.toHaveBeenCalled() - } finally { - vi.unstubAllGlobals() - } - }) - - it('should still track when clearing the pending entry fails', () => { - const pending = { eventName: 'user_registration_success', properties: { method: 'email' } } - vi.stubGlobal('window', { - sessionStorage: { - getItem: () => JSON.stringify(pending), - setItem: vi.fn(), removeItem: () => { throw new Error('remove failed') }, }, }) - try { - flushRegistrationSuccess() - - expect(mockTrackEvent).toHaveBeenCalledWith('user_registration_success', { - method: 'email', - }) - } finally { - vi.unstubAllGlobals() - } - }) - }) - - // Both producers and the consumer must degrade gracefully when sessionStorage is - // missing (SSR) or blocked (privacy mode / disabled storage). - describe('when sessionStorage is unavailable', () => { - it('should no-op without throwing when window is undefined', () => { - vi.stubGlobal('window', undefined) - - try { - expect(() => rememberRegistrationSuccess({ method: 'email' })).not.toThrow() - expect(() => flushRegistrationSuccess()).not.toThrow() - expect(mockTrackEvent).not.toHaveBeenCalled() - } finally { - vi.unstubAllGlobals() - } + await expect(flushRegistrationSuccess()).resolves.toBeUndefined() + expect(mockTrackEvent).not.toHaveBeenCalled() }) - it('should no-op without throwing when accessing sessionStorage throws', () => { + it('retains the same marker when acknowledgement removal fails', async () => { + rememberRegistrationSuccess({ method: 'email' }) + const raw = window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY) + const removeItem = vi.fn(() => { + throw new Error('remove failed') + }) + vi.stubGlobal('window', { + sessionStorage: { + getItem: () => raw, + setItem: vi.fn(), + removeItem, + }, + }) + + await expect(flushRegistrationSuccess()).resolves.toBeUndefined() + await expect(flushRegistrationSuccess()).resolves.toBeUndefined() + + expect(mockTrackEvent).toHaveBeenCalledTimes(2) + expect(mockTrackEvent.mock.calls[0]?.[2]).toEqual(mockTrackEvent.mock.calls[1]?.[2]) + expect(removeItem).toHaveBeenCalledTimes(2) + }) + + it('no-ops when sessionStorage access is blocked', async () => { vi.stubGlobal('window', { get sessionStorage() { throw new Error('storage disabled') }, }) - try { - expect(() => rememberRegistrationSuccess({ method: 'oauth' })).not.toThrow() - expect(() => flushRegistrationSuccess()).not.toThrow() - expect(mockTrackEvent).not.toHaveBeenCalled() - } finally { - vi.unstubAllGlobals() - } + expect(() => rememberRegistrationSuccess({ method: 'email' })).not.toThrow() + await expect(flushRegistrationSuccess()).resolves.toBeUndefined() + expect(mockTrackEvent).not.toHaveBeenCalled() }) }) }) diff --git a/web/app/components/base/amplitude/registration-consent-coordinator.tsx b/web/app/components/base/amplitude/registration-consent-coordinator.tsx new file mode 100644 index 00000000000..8d19d05cdc2 --- /dev/null +++ b/web/app/components/base/amplitude/registration-consent-coordinator.tsx @@ -0,0 +1,15 @@ +'use client' + +import { useEffect } from 'react' +import { useAnalyticsConsent } from '@/app/components/base/analytics-consent/consent-store' +import { coordinateRegistrationConsent } from './registration-tracking' + +export function RegistrationConsentCoordinator() { + const consent = useAnalyticsConsent() + + useEffect(() => { + coordinateRegistrationConsent(consent) + }, [consent]) + + return null +} diff --git a/web/app/components/base/amplitude/registration-session-state.ts b/web/app/components/base/amplitude/registration-session-state.ts new file mode 100644 index 00000000000..5eae9390cb1 --- /dev/null +++ b/web/app/components/base/amplitude/registration-session-state.ts @@ -0,0 +1,99 @@ +export const REGISTRATION_SUCCESS_STORAGE_KEY = 'pending_registration_success_event' +export const OAUTH_REGISTRATION_GA_SENT_KEY = 'oauth_registration_ga_sent' +const FLUSH_RETRY_DELAYS_MS = [1000, 4000, 16000] as const + +export const REGISTRATION_METHODS = ['email', 'oauth'] as const + +export const ATTRIBUTION_KEYS = [ + 'utm_source', + 'utm_medium', + 'utm_campaign', + 'utm_content', + 'utm_term', + 'slug', +] as const + +export type RegistrationMethod = (typeof REGISTRATION_METHODS)[number] +export type RegistrationAttribution = Partial> + +export type RegistrationIntent = { + registrationId: string + occurredAt: number + method: RegistrationMethod + attribution: RegistrationAttribution +} + +let registrationDeliveryGeneration = 0 +let flushRetryTimer: ReturnType | null = null +let flushRetryAttempt = 0 + +export const getRegistrationSessionStorage = (): Storage | null => { + try { + if (typeof window === 'undefined') return null + return window.sessionStorage + } catch { + return null + } +} + +export const removeStoredRegistrationMarker = (storage = getRegistrationSessionStorage()) => { + try { + storage?.removeItem(REGISTRATION_SUCCESS_STORAGE_KEY) + } catch {} +} + +export const hasSentOAuthRegistrationGA = () => { + try { + return getRegistrationSessionStorage()?.getItem(OAUTH_REGISTRATION_GA_SENT_KEY) === 'true' + } catch { + return false + } +} + +export const markOAuthRegistrationGASent = () => { + try { + getRegistrationSessionStorage()?.setItem(OAUTH_REGISTRATION_GA_SENT_KEY, 'true') + } catch {} +} + +export const clearOAuthRegistrationGAGuard = () => { + try { + getRegistrationSessionStorage()?.removeItem(OAUTH_REGISTRATION_GA_SENT_KEY) + } catch {} +} + +export const getRegistrationDeliveryGeneration = () => registrationDeliveryGeneration + +export const clearRegistrationFlushRetry = () => { + if (flushRetryTimer !== null) { + clearTimeout(flushRetryTimer) + flushRetryTimer = null + } + flushRetryAttempt = 0 +} + +export const scheduleRegistrationFlushRetry = (runFlush: () => void) => { + if (flushRetryAttempt >= FLUSH_RETRY_DELAYS_MS.length) return + + const delay = FLUSH_RETRY_DELAYS_MS[flushRetryAttempt] + flushRetryAttempt += 1 + const generation = registrationDeliveryGeneration + if (flushRetryTimer !== null) clearTimeout(flushRetryTimer) + + flushRetryTimer = setTimeout(() => { + flushRetryTimer = null + if (generation !== registrationDeliveryGeneration) return + runFlush() + }, delay) +} + +export const invalidateRegistrationDeliveryState = () => { + registrationDeliveryGeneration += 1 + clearRegistrationFlushRetry() + removeStoredRegistrationMarker() +} + +export const discardRegistrationSessionState = () => { + invalidateRegistrationDeliveryState() + clearOAuthRegistrationGAGuard() +} diff --git a/web/app/components/base/amplitude/registration-tracking.ts b/web/app/components/base/amplitude/registration-tracking.ts index 5562d2173c4..15c11bc6590 100644 --- a/web/app/components/base/amplitude/registration-tracking.ts +++ b/web/app/components/base/amplitude/registration-tracking.ts @@ -1,38 +1,121 @@ +import type { + RegistrationAttribution, + RegistrationIntent, + RegistrationMethod, +} from './registration-session-state' +import type { AnalyticsConsent } from '@/app/components/base/analytics-consent/consent-store' import { getAnalyticsConsent } from '@/app/components/base/analytics-consent/consent-store' +import { getIsAmplitudeInitialized } from './init' +import { + ATTRIBUTION_KEYS, + clearRegistrationFlushRetry, + getRegistrationDeliveryGeneration, + getRegistrationSessionStorage, + invalidateRegistrationDeliveryState, + REGISTRATION_METHODS, + REGISTRATION_SUCCESS_STORAGE_KEY, + removeStoredRegistrationMarker, + scheduleRegistrationFlushRetry, +} from './registration-session-state' import { trackEvent } from './utils' -/** - * Storage key for a registration success event that is waiting to be sent to - * Amplitude until a user ID has been attached. - */ -export const REGISTRATION_SUCCESS_STORAGE_KEY = 'pending_registration_success_event' +const REGISTRATION_MARKER_VERSION = 2 +const REGISTRATION_MARKER_TTL_MS = 24 * 60 * 60 * 1000 +// Browser clocks may be corrected between registration and delivery. Permit a small +// correction, but reject timestamps far enough ahead to corrupt Amplitude ordering. +const REGISTRATION_FUTURE_CLOCK_SKEW_ALLOWANCE_MS = 5 * 60 * 1000 +const SUCCESSFUL_TRACK_RESULT_MIN = 200 +const SUCCESSFUL_TRACK_RESULT_MAX = 299 -type RegistrationMethod = 'email' | 'oauth' +const REGISTRATION_EVENT_NAMES = [ + 'user_registration_success', + 'user_registration_success_with_utm', +] as const -type PendingRegistrationSuccessEvent = { - eventName: string - properties: Record +type RegistrationEventName = (typeof REGISTRATION_EVENT_NAMES)[number] + +type PendingRegistrationSuccessEvent = RegistrationIntent & { + version: typeof REGISTRATION_MARKER_VERSION + expiresAt: number + eventName: RegistrationEventName } -const getSessionStorage = (): Storage | null => { +let registrationSnapshot = 0 +let activeFlush: { generation: number; promise: Promise } | null = null +const registrationListeners = new Set<() => void>() + +const isRecord = (value: unknown): value is Record => + Boolean(value) && typeof value === 'object' && !Array.isArray(value) + +const isRegistrationMethod = (value: unknown): value is RegistrationMethod => + typeof value === 'string' && REGISTRATION_METHODS.includes(value as RegistrationMethod) + +const isRegistrationEventName = (value: unknown): value is RegistrationEventName => + typeof value === 'string' && REGISTRATION_EVENT_NAMES.includes(value as RegistrationEventName) + +const notifyRegistrationMarkerStored = () => { + registrationSnapshot += 1 + registrationListeners.forEach((listener) => listener()) +} + +const createRegistrationId = () => { try { - if (typeof window === 'undefined') return null - return window.sessionStorage + return globalThis.crypto.randomUUID() } catch { - return null + return `${Date.now()}-${Math.random().toString(36).slice(2)}` + } +} + +export const normalizeRegistrationAttribution = ( + value?: Record | null, +): RegistrationAttribution | null => { + if (!value) return null + + const attribution: RegistrationAttribution = {} + ATTRIBUTION_KEYS.forEach((key) => { + const item = value[key] + if (typeof item !== 'string') return + + const normalized = item.trim() + if (normalized) attribution[key] = normalized + }) + + return Object.keys(attribution).length ? attribution : null +} + +const createRegistrationIntent = ( + method: RegistrationMethod, + utmInfo?: Record | null, +): RegistrationIntent => ({ + registrationId: createRegistrationId(), + occurredAt: Date.now(), + method, + attribution: normalizeRegistrationAttribution(utmInfo) ?? {}, +}) + +const storeRegistrationIntent = (intent: RegistrationIntent) => { + const storage = getRegistrationSessionStorage() + if (!storage) return false + + const pending: PendingRegistrationSuccessEvent = { + ...intent, + version: REGISTRATION_MARKER_VERSION, + expiresAt: intent.occurredAt + REGISTRATION_MARKER_TTL_MS, + eventName: Object.keys(intent.attribution).length + ? 'user_registration_success_with_utm' + : 'user_registration_success', + } + + try { + storage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, JSON.stringify(pending)) + clearRegistrationFlushRetry() + notifyRegistrationMarkerStored() + return true + } catch { + return false } } -/** - * Remember a registration success event after analytics consent so it can be sent - * to Amplitude *after* the user ID is attached (see `flushRegistrationSuccess`). - * - * Amplitude attributes events to whatever identity is active when `track` runs. At - * registration time the client does not yet know the user ID, so firing the event - * immediately records it under an anonymous profile. We persist the event here and - * replay it once `setUserId` runs in the bootstrap effects after the redirect. An - * event produced before analytics consent is granted is dropped instead of queued. - */ export const rememberRegistrationSuccess = ({ method, utmInfo, @@ -40,49 +123,157 @@ export const rememberRegistrationSuccess = ({ method: RegistrationMethod utmInfo?: Record | null }) => { - if (getAnalyticsConsent() !== 'granted') return - - const storage = getSessionStorage() - if (!storage) return - - const pending: PendingRegistrationSuccessEvent = { - eventName: utmInfo ? 'user_registration_success_with_utm' : 'user_registration_success', - properties: { method, ...utmInfo }, + const consent = getAnalyticsConsent() + if (consent === 'denied' || consent === 'disabled') { + invalidateRegistrationDeliveryState() + return false } - try { - storage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, JSON.stringify(pending)) - } catch {} + // Persist even while consent is unknown. Flush waits for grant + Amplitude init + // + user identity, so a later full-page redirect still has the marker. + return storeRegistrationIntent(createRegistrationIntent(method, utmInfo)) } -/** - * Send a previously remembered registration success event to Amplitude. - * - * MUST be called after `setUserId` so the event lands on the identified user profile. - * No-op when nothing is pending. The pending entry is removed before tracking so the - * event fires at most once even if this runs multiple times. - */ -export const flushRegistrationSuccess = () => { - const storage = getSessionStorage() +export const coordinateRegistrationConsent = (consent: AnalyticsConsent) => { + if (consent === 'denied' || consent === 'disabled') invalidateRegistrationDeliveryState() +} + +export const subscribeRegistrationSuccess = (listener: () => void) => { + registrationListeners.add(listener) + return () => registrationListeners.delete(listener) +} + +export const getRegistrationSuccessSnapshot = () => registrationSnapshot + +const isRegistrationAttribution = (value: unknown): value is RegistrationAttribution => { + if (!isRecord(value)) return false + + return Object.entries(value).every( + ([key, item]) => + ATTRIBUTION_KEYS.includes(key as (typeof ATTRIBUTION_KEYS)[number]) && + typeof item === 'string' && + Boolean(item.trim()), + ) +} + +const parsePendingRegistration = (raw: string): PendingRegistrationSuccessEvent | null => { + try { + const value: unknown = JSON.parse(raw) + if (!isRecord(value)) return null + if (value.version !== REGISTRATION_MARKER_VERSION) return null + if (typeof value.registrationId !== 'string' || !value.registrationId) return null + if (typeof value.occurredAt !== 'number' || !Number.isFinite(value.occurredAt)) return null + if (typeof value.expiresAt !== 'number' || !Number.isFinite(value.expiresAt)) return null + if (value.expiresAt !== value.occurredAt + REGISTRATION_MARKER_TTL_MS) return null + if (!isRegistrationEventName(value.eventName)) return null + if (!isRegistrationMethod(value.method)) return null + if (!isRegistrationAttribution(value.attribution)) return null + + const hasAttribution = Object.keys(value.attribution).length > 0 + if (hasAttribution !== (value.eventName === 'user_registration_success_with_utm')) return null + + return value as PendingRegistrationSuccessEvent + } catch { + return null + } +} + +const runRegistrationFlush = async (generation: number) => { + const isStale = () => generation !== getRegistrationDeliveryGeneration() + if (isStale()) return + + const consent = getAnalyticsConsent() + if (consent === 'unknown') return + + const storage = getRegistrationSessionStorage() if (!storage) return - let raw: string | null = null - try { - raw = storage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY) - } catch { + if (consent === 'denied' || consent === 'disabled') { + invalidateRegistrationDeliveryState() return } + if (!getIsAmplitudeInitialized()) return - if (!raw) return + while (true) { + if (isStale()) return - try { - storage.removeItem(REGISTRATION_SUCCESS_STORAGE_KEY) - } catch {} + let raw: string | null + try { + raw = storage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY) + } catch { + return + } + if (!raw) return - if (getAnalyticsConsent() !== 'granted') return + const pending = parsePendingRegistration(raw) + const now = Date.now() + if ( + !pending || + pending.expiresAt <= now || + pending.occurredAt > now + REGISTRATION_FUTURE_CLOCK_SKEW_ALLOWANCE_MS + ) { + removeStoredRegistrationMarker(storage) + return + } - try { - const pending = JSON.parse(raw) as PendingRegistrationSuccessEvent - if (pending?.eventName) trackEvent(pending.eventName, pending.properties) - } catch {} + let trackResult: ReturnType + try { + trackResult = trackEvent( + pending.eventName, + { + method: pending.method, + ...pending.attribution, + registration_id: pending.registrationId, + event_version: REGISTRATION_MARKER_VERSION, + tracking_contract_version: 'consent_wait_v2', + }, + { + insert_id: pending.registrationId, + time: pending.occurredAt, + }, + ) + } catch { + return + } + if (!trackResult) return + + let acknowledged = false + try { + const result: { code?: unknown } = await trackResult.promise + acknowledged = + typeof result.code === 'number' && + result.code >= SUCCESSFUL_TRACK_RESULT_MIN && + result.code <= SUCCESSFUL_TRACK_RESULT_MAX + } catch {} + if (isStale()) return + + let currentRaw: string | null + try { + currentRaw = storage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY) + } catch { + return + } + if (currentRaw !== raw) continue + if (!acknowledged) { + scheduleRegistrationFlushRetry(() => { + void flushRegistrationSuccess() + }) + return + } + + clearRegistrationFlushRetry() + removeStoredRegistrationMarker(storage) + return + } +} + +export function flushRegistrationSuccess() { + const generation = getRegistrationDeliveryGeneration() + if (activeFlush?.generation === generation) return activeFlush.promise + + const promise = runRegistrationFlush(generation).finally(() => { + if (activeFlush?.promise === promise) activeFlush = null + }) + activeFlush = { generation, promise } + return promise } diff --git a/web/app/components/base/amplitude/utils.ts b/web/app/components/base/amplitude/utils.ts index 58354463fc1..bb6021d0ad1 100644 --- a/web/app/components/base/amplitude/utils.ts +++ b/web/app/components/base/amplitude/utils.ts @@ -9,8 +9,13 @@ const canUseAmplitude = () => getAnalyticsConsent() === 'granted' && getIsAmplit * @param eventName Event name * @param eventProperties Event properties (optional) */ -export const trackEvent = (eventName: string, eventProperties?: Record) => { +export const trackEvent = ( + eventName: string, + eventProperties?: Record, + eventOptions?: amplitude.Types.EventOptions, +) => { if (!canUseAmplitude()) return + if (eventOptions) return amplitude.track(eventName, eventProperties, eventOptions) return amplitude.track(eventName, eventProperties) } diff --git a/web/app/components/base/analytics-consent/__tests__/analytics-disabled.spec.tsx b/web/app/components/base/analytics-consent/__tests__/analytics-disabled.spec.tsx new file mode 100644 index 00000000000..65443ba0046 --- /dev/null +++ b/web/app/components/base/analytics-consent/__tests__/analytics-disabled.spec.tsx @@ -0,0 +1,26 @@ +import { render, waitFor } from '@testing-library/react' +import { REGISTRATION_SUCCESS_STORAGE_KEY } from '../../amplitude/registration-session-state' +import { + coordinateRegistrationConsent, + rememberRegistrationSuccess, +} from '../../amplitude/registration-tracking' +import { AnalyticsDisabled } from '../analytics-disabled' +import { getAnalyticsConsent, setAnalyticsConsent } from '../consent-store' + +describe('AnalyticsDisabled', () => { + beforeEach(() => { + window.sessionStorage.clear() + coordinateRegistrationConsent('denied') + setAnalyticsConsent('granted') + }) + + it('terminally discards a pending registration marker', async () => { + rememberRegistrationSuccess({ method: 'email' }) + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).not.toBeNull() + + render() + + await waitFor(() => expect(getAnalyticsConsent()).toBe('disabled')) + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + }) +}) diff --git a/web/app/components/base/analytics-consent/__tests__/analytics-runtimes.spec.tsx b/web/app/components/base/analytics-consent/__tests__/analytics-runtimes.spec.tsx index bfad2dc720a..d7dabf25264 100644 --- a/web/app/components/base/analytics-consent/__tests__/analytics-runtimes.spec.tsx +++ b/web/app/components/base/analytics-consent/__tests__/analytics-runtimes.spec.tsx @@ -14,6 +14,10 @@ vi.mock('@/app/components/base/amplitude/WebAppAmplitudeProvider', () => ({ WebAppAmplitudeProvider: () => , })) +vi.mock('@/app/components/base/amplitude/registration-consent-coordinator', () => ({ + RegistrationConsentCoordinator: () => , +})) + vi.mock('@/app/components/external-attribution-recorder', () => ({ default: () => , })) @@ -24,6 +28,7 @@ describe('analytics runtimes', () => { expect(screen.getByTestId('cookieyes-consent-bridge')).toBeInTheDocument() expect(screen.getByTestId('console-amplitude-provider')).toBeInTheDocument() + expect(screen.getByTestId('registration-consent-coordinator')).toBeInTheDocument() expect(screen.getByTestId('external-attribution-recorder')).toBeInTheDocument() expect(screen.queryByTestId('web-app-amplitude-provider')).toBeNull() }) @@ -34,6 +39,7 @@ describe('analytics runtimes', () => { expect(screen.getByTestId('cookieyes-consent-bridge')).toBeInTheDocument() expect(screen.getByTestId('web-app-amplitude-provider')).toBeInTheDocument() expect(screen.queryByTestId('console-amplitude-provider')).toBeNull() + expect(screen.queryByTestId('registration-consent-coordinator')).toBeNull() expect(screen.queryByTestId('external-attribution-recorder')).toBeNull() }) }) diff --git a/web/app/components/base/analytics-consent/__tests__/cloud-analytics.spec.tsx b/web/app/components/base/analytics-consent/__tests__/cloud-analytics.spec.tsx index 2380b7485c0..a4dc0dc299b 100644 --- a/web/app/components/base/analytics-consent/__tests__/cloud-analytics.spec.tsx +++ b/web/app/components/base/analytics-consent/__tests__/cloud-analytics.spec.tsx @@ -1,5 +1,6 @@ import { QueryClient } from '@tanstack/react-query' import { render } from '@testing-library/react' +import { REGISTRATION_SUCCESS_STORAGE_KEY } from '@/app/components/base/amplitude/registration-session-state' let queryClient: QueryClient @@ -55,9 +56,13 @@ vi.mock('../cloud-analytics-layout-boundary', () => ({ ), })) -async function renderCloudAnalytics() { +async function getCloudAnalyticsResult() { const { CloudAnalytics } = await import('../cloud-analytics') - return render(await CloudAnalytics()) + return CloudAnalytics() +} + +async function renderCloudAnalytics() { + return render(await getCloudAnalyticsResult()) } describe('CloudAnalytics', () => { @@ -68,6 +73,7 @@ describe('CloudAnalytics', () => { configState.isProd = true configState.webPrefix = 'https://cloud.dify.ai' queryClient = new QueryClient() + window.sessionStorage.clear() queryClient.setQueryData(systemFeaturesQueryKey, { deployment_edition: 'CLOUD' }) mockHeadersGet.mockImplementation((name: string) => { const values: Record = { @@ -94,9 +100,13 @@ describe('CloudAnalytics', () => { return values[name] ?? null }) - const { queryByTestId } = await renderCloudAnalytics() + const result = await getCloudAnalyticsResult() + const { queryByTestId } = render(result) expect(queryByTestId('cloud-analytics-layout-boundary')).toBeNull() + expect(result).not.toBeNull() + const { getAnalyticsConsent } = await import('../consent-store') + expect(getAnalyticsConsent()).toBe('disabled') }) it.each(['COMMUNITY', 'ENTERPRISE'] as const)( @@ -109,11 +119,16 @@ describe('CloudAnalytics', () => { }, ) - it('does not render when System Features are unavailable', async () => { + it('suspends analytics without deleting pending registration state when System Features are unavailable', async () => { queryClient.removeQueries({ queryKey: systemFeaturesQueryKey }) + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, 'pending-marker') - const { queryByTestId } = await renderCloudAnalytics() + const result = await getCloudAnalyticsResult() + const { queryByTestId } = render(result) expect(queryByTestId('cloud-analytics-layout-boundary')).toBeNull() + const { getAnalyticsConsent } = await import('../consent-store') + expect(getAnalyticsConsent()).toBe('unknown') + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBe('pending-marker') }) }) diff --git a/web/app/components/base/analytics-consent/analytics-disabled.tsx b/web/app/components/base/analytics-consent/analytics-disabled.tsx new file mode 100644 index 00000000000..a562e67c51a --- /dev/null +++ b/web/app/components/base/analytics-consent/analytics-disabled.tsx @@ -0,0 +1,14 @@ +'use client' + +import { useEffect } from 'react' +import { coordinateRegistrationConsent } from '@/app/components/base/amplitude/registration-tracking' +import { setAnalyticsConsent } from './consent-store' + +export function AnalyticsDisabled() { + useEffect(() => { + setAnalyticsConsent('disabled') + coordinateRegistrationConsent('disabled') + }, []) + + return null +} diff --git a/web/app/components/base/analytics-consent/cloud-analytics.tsx b/web/app/components/base/analytics-consent/cloud-analytics.tsx index 9b53224b494..cf4bd0055f7 100644 --- a/web/app/components/base/analytics-consent/cloud-analytics.tsx +++ b/web/app/components/base/analytics-consent/cloud-analytics.tsx @@ -1,6 +1,7 @@ import { COOKIEYES_SITE_KEY, IS_PROD, WEB_PREFIX } from '@/config' import { getCachedSystemFeatures } from '@/features/system-features/server' import { headers } from '@/next/headers' +import { AnalyticsDisabled } from './analytics-disabled' import { CloudAnalyticsLayoutBoundary } from './cloud-analytics-layout-boundary' import { isCloudAnalyticsRequest } from './request-boundary' @@ -19,7 +20,7 @@ export async function CloudAnalytics() { webPrefix: WEB_PREFIX, }) - if (!enabled) return null + if (!enabled) return const nonce = requestHeaders.get('x-nonce') ?? undefined diff --git a/web/app/components/base/analytics-consent/consent-store.ts b/web/app/components/base/analytics-consent/consent-store.ts index 030d3f770b7..b1f77c2eac7 100644 --- a/web/app/components/base/analytics-consent/consent-store.ts +++ b/web/app/components/base/analytics-consent/consent-store.ts @@ -2,7 +2,7 @@ import { useSyncExternalStore } from 'react' -export type AnalyticsConsent = 'unknown' | 'denied' | 'granted' +export type AnalyticsConsent = 'unknown' | 'denied' | 'granted' | 'disabled' type CookieYesConsentUpdateDetail = { accepted: string[] diff --git a/web/app/components/base/analytics-consent/console-analytics-runtime.tsx b/web/app/components/base/analytics-consent/console-analytics-runtime.tsx index b0f97e5cd7e..00814fad51c 100644 --- a/web/app/components/base/analytics-consent/console-analytics-runtime.tsx +++ b/web/app/components/base/analytics-consent/console-analytics-runtime.tsx @@ -1,6 +1,7 @@ 'use client' import AmplitudeProvider from '@/app/components/base/amplitude' +import { RegistrationConsentCoordinator } from '@/app/components/base/amplitude/registration-consent-coordinator' import ExternalAttributionRecorder from '@/app/components/external-attribution-recorder' import { CookieYesConsentBridge } from './cookieyes-consent-bridge' @@ -9,6 +10,7 @@ export function ConsoleAnalyticsRuntime() { <> + ) diff --git a/web/app/components/oauth-registration-analytics.tsx b/web/app/components/oauth-registration-analytics.tsx index fd7eb3bc542..e82adaf2923 100644 --- a/web/app/components/oauth-registration-analytics.tsx +++ b/web/app/components/oauth-registration-analytics.tsx @@ -2,9 +2,18 @@ import Cookies from 'js-cookie' import { useEffect, useRef } from 'react' +import { useAnalyticsConsent } from '@/app/components/base/analytics-consent/consent-store' import { useSearchParams } from '@/next/navigation' import { sendGAEvent } from '@/utils/gtag' -import { rememberRegistrationSuccess } from './base/amplitude/registration-tracking' +import { + clearOAuthRegistrationGAGuard, + hasSentOAuthRegistrationGA, + markOAuthRegistrationGASent, +} from './base/amplitude/registration-session-state' +import { + normalizeRegistrationAttribution, + rememberRegistrationSuccess, +} from './base/amplitude/registration-tracking' const OAUTH_NEW_USER_PARAM = 'oauth_new_user' @@ -18,46 +27,75 @@ const removeOAuthNewUserParam = () => { } export function OAuthRegistrationAnalytics() { + const analyticsConsent = useAnalyticsConsent() const searchParams = useSearchParams() const oauthNewUserParam = searchParams.get(OAUTH_NEW_USER_PARAM) - const handledParamRef = useRef(null) + const gaHandledRef = useRef(false) + const amplitudeHandledRef = useRef(false) + const cleanedRef = useRef(false) + const utmInfoRef = useRef | undefined>( + undefined, + ) useEffect(() => { - if (oauthNewUserParam === null || handledParamRef.current === oauthNewUserParam) return - - handledParamRef.current = oauthNewUserParam - const oauthNewUser = oauthNewUserParam === 'true' - if (!oauthNewUser) { - removeOAuthNewUserParam() + if (oauthNewUserParam === null) { + clearOAuthRegistrationGAGuard() return } - let utmInfo: Record | null = null - const utmInfoStr = Cookies.get('utm_info') - if (utmInfoStr) { - try { - const parsed: unknown = JSON.parse(utmInfoStr) - if (isRecord(parsed)) utmInfo = parsed - } catch (e) { - console.error('Failed to parse utm_info cookie:', e) + const oauthNewUser = oauthNewUserParam === 'true' + if (!oauthNewUser) { + if (!cleanedRef.current) { + cleanedRef.current = true + clearOAuthRegistrationGAGuard() + removeOAuthNewUserParam() } + return } + if (utmInfoRef.current === undefined) { + let parsedUtmInfo: Record | null = null + const utmInfoStr = Cookies.get('utm_info') + if (utmInfoStr) { + try { + const parsed: unknown = JSON.parse(utmInfoStr) + if (isRecord(parsed)) parsedUtmInfo = parsed + } catch (e) { + console.error('Failed to parse utm_info cookie:', e) + } + } + utmInfoRef.current = normalizeRegistrationAttribution(parsedUtmInfo) + } + const utmInfo = utmInfoRef.current + const eventName = utmInfo ? 'user_registration_success_with_utm' : 'user_registration_success' - // Defer the Amplitude event until the user ID is attached. The app context - // external sync replays it after setUserId runs. Firing it here would record it under an - // anonymous Amplitude profile (no user ID set yet). - rememberRegistrationSuccess({ method: 'oauth', utmInfo }) + if (!gaHandledRef.current) { + gaHandledRef.current = true + if (!hasSentOAuthRegistrationGA()) { + sendGAEvent(eventName, { + method: 'oauth', + ...utmInfo, + }) + markOAuthRegistrationGASent() + } + } - sendGAEvent(eventName, { - method: 'oauth', - ...utmInfo, - }) + if ( + (analyticsConsent === 'unknown' || analyticsConsent === 'granted') && + !amplitudeHandledRef.current + ) { + const persisted = rememberRegistrationSuccess({ method: 'oauth', utmInfo }) + if (!persisted) return + amplitudeHandledRef.current = true + } - Cookies.remove('utm_info') - removeOAuthNewUserParam() - }, [oauthNewUserParam]) + if (!cleanedRef.current) { + cleanedRef.current = true + Cookies.remove('utm_info') + removeOAuthNewUserParam() + } + }, [analyticsConsent, oauthNewUserParam]) return null } diff --git a/web/context/__tests__/console-bootstrap.spec.tsx b/web/context/__tests__/console-bootstrap.spec.tsx index 9e0e66bbbb1..3287a77cd5a 100644 --- a/web/context/__tests__/console-bootstrap.spec.tsx +++ b/web/context/__tests__/console-bootstrap.spec.tsx @@ -186,6 +186,8 @@ vi.mock('@/app/components/base/amplitude/use-amplitude-initialized', () => ({ vi.mock('@/app/components/base/amplitude/registration-tracking', () => ({ flushRegistrationSuccess: vi.fn(), + subscribeRegistrationSuccess: () => () => {}, + getRegistrationSuccessSnapshot: () => 0, })) vi.mock('@/app/components/base/zendesk/utils', () => ({ diff --git a/web/service/__tests__/base-request.spec.ts b/web/service/__tests__/base-request.spec.ts index c324a83cb55..0da529bc10c 100644 --- a/web/service/__tests__/base-request.spec.ts +++ b/web/service/__tests__/base-request.spec.ts @@ -1,4 +1,9 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vite-plus/test' +import { + discardRegistrationSessionState, + OAUTH_REGISTRATION_GA_SENT_KEY, + REGISTRATION_SUCCESS_STORAGE_KEY, +} from '@/app/components/base/amplitude/registration-session-state' // oxlint-disable-next-line no-restricted-imports -- This spec directly tests the legacy request owner. import { request } from '../base' @@ -51,6 +56,21 @@ const createUnauthorizedResponse = () => }, ) +const createForcedLogoutResponse = () => + new Response( + JSON.stringify({ + code: 'unauthorized_and_force_logout', + message: 'This account session is no longer valid.', + status: 401, + }), + { + status: 401, + headers: { + 'Content-Type': 'application/json', + }, + }, + ) + type ClientRequestOptions = { response: Response refreshError?: Error @@ -80,9 +100,11 @@ describe('request 401 handling', () => { writable: true, configurable: true, }) + window.sessionStorage.clear() }) afterEach(() => { + discardRegistrationSessionState() Object.defineProperty(globalThis, 'location', { value: originalLocation, writable: true, @@ -103,21 +125,42 @@ describe('request 401 handling', () => { it('should preserve the current URL when a 401 response cannot be parsed', async () => { const response = new Response('not-json', { status: 401 }) arrangeClientRequest({ response }) + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, 'account-a-marker') await expect(request('/account/profile')).rejects.toBe(response) expect(globalThis.location.href).toBe( `https://example.com/app/signin?redirect_url=${encodeURIComponent('/app/apps?category=agent#recent')}`, ) + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() expect(mocks.refreshAccessTokenOrReLogin).not.toHaveBeenCalled() }) + it('clears account A registration state before a forced reload so account B starts clean', async () => { + const response = createForcedLogoutResponse() + arrangeClientRequest({ response }) + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, 'account-a-marker') + window.sessionStorage.setItem(OAUTH_REGISTRATION_GA_SENT_KEY, 'true') + const reload = vi.fn(() => { + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + expect(window.sessionStorage.getItem(OAUTH_REGISTRATION_GA_SENT_KEY)).toBeNull() + }) + globalThis.location.reload = reload + + await expect(request('/account/profile')).rejects.toBe(response) + + expect(reload).toHaveBeenCalledOnce() + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, 'account-b-marker') + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBe('account-b-marker') + }) + it('should preserve the current URL when token refresh fails', async () => { const response = createUnauthorizedResponse() arrangeClientRequest({ response, refreshError: new Error('refresh failed'), }) + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, 'account-a-marker') await expect(request('/account/profile')).rejects.toBe(response) @@ -125,5 +168,16 @@ describe('request 401 handling', () => { expect(globalThis.location.href).toBe( `https://example.com/app/signin?redirect_url=${encodeURIComponent('/app/apps?category=agent#recent')}`, ) + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + }) + + it('does not clear console registration state for a public-app 401 redirect', async () => { + const response = createUnauthorizedResponse() + arrangeClientRequest({ response }) + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, 'console-marker') + + await expect(request('/account/profile', {}, { isPublicAPI: true })).rejects.toBe(response) + + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBe('console-marker') }) }) diff --git a/web/service/base.ts b/web/service/base.ts index d5ece4cd7d4..ea2435fdb6b 100644 --- a/web/service/base.ts +++ b/web/service/base.ts @@ -30,6 +30,7 @@ import type { } from '@/types/workflow' import { toast } from '@langgenius/dify-ui/toast' import Cookies from 'js-cookie' +import { discardRegistrationSessionState } from '@/app/components/base/amplitude/registration-session-state' import { API_PREFIX, CSRF_COOKIE_NAME, @@ -197,6 +198,14 @@ export type IOtherOptions = { onDataSourceNodeError?: IOnDataSourceNodeError } +const discardRegistrationStateForConsoleAuthBoundary = ({ + isMarketplaceAPI, + isPublicAPI, +}: IOtherOptions) => { + if (isMarketplaceAPI || isPublicAPI) return + discardRegistrationSessionState() +} + function jumpTo(url: string) { if (!url || !isClient) return const targetPath = new URL(url, window.location.origin).pathname @@ -1008,6 +1017,7 @@ export const request = async (url: string, options = {}, otherOptions?: IOthe const [parseErr, errRespData] = await asyncRunSafe(errResp.json()) if (parseErr) { + discardRegistrationStateForConsoleAuthBoundary(otherOptionsForBaseFetch) window.location.href = buildSigninUrlWithRedirect() return Promise.reject(err) } @@ -1025,6 +1035,7 @@ export const request = async (url: string, options = {}, otherOptions?: IOthe } if (code === 'unauthorized_and_force_logout') { // Cookies will be cleared by the backend + discardRegistrationStateForConsoleAuthBoundary(otherOptionsForBaseFetch) window.location.reload() return Promise.reject(err) } @@ -1053,6 +1064,7 @@ export const request = async (url: string, options = {}, otherOptions?: IOthe // there. Redirecting to /signin loses the user_code context and // the post-login flow lands on /apps instead of returning here. if (window.location.pathname === `${basePath}/device`) return Promise.reject(err) + discardRegistrationStateForConsoleAuthBoundary(otherOptionsForBaseFetch) if (window.location.pathname !== `${basePath}/signin`) { jumpTo(buildSigninUrlWithRedirect()) return Promise.reject(err) diff --git a/web/service/common.spec.ts b/web/service/common.spec.ts index ef678cb69f9..cfbf09ad331 100644 --- a/web/service/common.spec.ts +++ b/web/service/common.spec.ts @@ -1,5 +1,14 @@ +import type { ReactNode } from 'react' +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { act, renderHook } from '@testing-library/react' +import { createElement } from 'react' import { beforeEach, describe, expect, it, vi } from 'vite-plus/test' +import { + OAUTH_REGISTRATION_GA_SENT_KEY, + REGISTRATION_SUCCESS_STORAGE_KEY, +} from '@/app/components/base/amplitude/registration-session-state' import { emailLoginWithCode, sendEMailLoginCode } from './common' +import { useLogout } from './use-common' const mocks = vi.hoisted(() => ({ post: vi.fn(), @@ -68,3 +77,29 @@ describe('emailLoginWithCode', () => { }) }) }) + +describe('useLogout', () => { + beforeEach(() => { + vi.clearAllMocks() + window.sessionStorage.clear() + }) + + it('discards registration delivery state after a successful logout', async () => { + const queryClient = new QueryClient() + queryClient.setQueryData(['account-profile'], { id: 'previous-user' }) + window.sessionStorage.setItem(REGISTRATION_SUCCESS_STORAGE_KEY, 'pending-marker') + window.sessionStorage.setItem(OAUTH_REGISTRATION_GA_SENT_KEY, 'true') + mocks.post.mockResolvedValueOnce({ result: 'success' }) + const wrapper = ({ children }: { children: ReactNode }) => + createElement(QueryClientProvider, { client: queryClient }, children) + const { result } = renderHook(() => useLogout(), { wrapper }) + + await act(async () => { + await result.current.mutateAsync() + }) + + expect(window.sessionStorage.getItem(REGISTRATION_SUCCESS_STORAGE_KEY)).toBeNull() + expect(window.sessionStorage.getItem(OAUTH_REGISTRATION_GA_SENT_KEY)).toBeNull() + expect(queryClient.getQueryData(['account-profile'])).toBeUndefined() + }) +}) diff --git a/web/service/use-common.ts b/web/service/use-common.ts index 95564c5fbea..524a38bfd0f 100644 --- a/web/service/use-common.ts +++ b/web/service/use-common.ts @@ -17,6 +17,7 @@ import type { } from '@/models/common' import type { RETRIEVE_METHOD } from '@/types/app' import { queryOptions, useMutation, useQuery, useQueryClient } from '@tanstack/react-query' +import { discardRegistrationSessionState } from '@/app/components/base/amplitude/registration-session-state' // oxlint-disable-next-line no-restricted-imports import { get, post } from './base' import { consoleQuery } from './client' @@ -162,6 +163,7 @@ export const useLogout = () => { mutationKey: [NAME_SPACE, 'logout'], mutationFn: () => post('/logout'), onSuccess: () => { + discardRegistrationSessionState() // Drop all cached queries so the post-logout /signin probe doesn't read // the previous user's profile (the userProfile queryKey is shared with // the (commonLayout) tree, which keeps observing it during React's From d25e2b51fcffe7629ecd6bb8ece2d385ec2e6122 Mon Sep 17 00:00:00 2001 From: "Byron.wang" Date: Mon, 31 Aug 2026 04:37:07 +0000 Subject: [PATCH 08/21] refactor(api): decouple onboarding and notification services (#40759) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: hjlarry --- api/controllers/console/notification.py | 104 +---- api/controllers/console/onboarding.py | 57 +-- api/dev/generate_swagger_markdown_docs.py | 4 + api/extensions/ext_application_services.py | 16 + api/openapi/markdown/console-openapi.md | 50 +-- api/openapi/markdown/service-openapi.md | 34 +- api/openapi/markdown/web-openapi.md | 6 +- .../step_by_step_tour_repository.py | 189 +++++++++ .../entities/notification_entities.py | 38 ++ api/services/entities/onboarding_entities.py | 42 ++ api/services/notification_gateway.py | 48 +++ api/services/notification_service.py | 60 +++ api/services/step_by_step_tour_service.py | 296 ++++++-------- .../test_generate_swagger_markdown_docs.py | 36 ++ .../controllers/console/test_notification.py | 77 ++++ .../controllers/console/test_onboarding.py | 102 ++--- .../test_ext_application_services.py | 2 + .../test_step_by_step_tour_repository.py | 171 ++++++++ .../services/test_notification_gateway.py | 63 +++ .../services/test_notification_service.py | 138 +++++++ .../test_step_by_step_tour_service.py | 369 +++++++++--------- 21 files changed, 1305 insertions(+), 597 deletions(-) create mode 100644 api/repositories/step_by_step_tour_repository.py create mode 100644 api/services/entities/notification_entities.py create mode 100644 api/services/entities/onboarding_entities.py create mode 100644 api/services/notification_gateway.py create mode 100644 api/services/notification_service.py create mode 100644 api/tests/unit_tests/controllers/console/test_notification.py create mode 100644 api/tests/unit_tests/repositories/test_step_by_step_tour_repository.py create mode 100644 api/tests/unit_tests/services/test_notification_gateway.py create mode 100644 api/tests/unit_tests/services/test_notification_service.py diff --git a/api/controllers/console/notification.py b/api/controllers/console/notification.py index 3e58f598bf7..080080bb361 100644 --- a/api/controllers/console/notification.py +++ b/api/controllers/console/notification.py @@ -1,56 +1,16 @@ -from collections.abc import Mapping -from typing import TypedDict - from flask_restx import Resource from pydantic import BaseModel, Field from controllers.common.fields import SimpleResultResponse from controllers.common.schema import register_response_schema_models, register_schema_models from controllers.console import console_ns -from controllers.console.wraps import ( - account_initialization_required, - model_validate, - only_edition_cloud, - setup_required, - with_current_user, -) +from controllers.console.flask_admission import console_account_admission +from controllers.console.wraps import model_validate +from enums import DeploymentEdition +from extensions.ext_application_services import application_services from fields.base import ResponseModel -from libs.login import login_required -from models import Account -from services.billing_service import BillingService - -# Notification content is stored under three lang tags. -_FALLBACK_LANG = "en-US" - - -class NotificationLangContent(TypedDict, total=False): - lang: str - title: str - subtitle: str - body: str - titlePicUrl: str - - -class NotificationItemDict(TypedDict): - notification_id: str | None - frequency: str | None - lang: str - title: str - subtitle: str - body: str - title_pic_url: str - - -class NotificationResponseDict(TypedDict): - should_show: bool - notifications: list[NotificationItemDict] - - -def _pick_lang_content(contents: Mapping[str, NotificationLangContent], lang: str) -> NotificationLangContent: - """Return the single LangContent for *lang*, falling back to English.""" - return ( - contents.get(lang) or contents.get(_FALLBACK_LANG) or next(iter(contents.values()), NotificationLangContent()) - ) +from libs.helper import dump_response +from machinery.context import RequestContext class DismissNotificationPayload(BaseModel): @@ -92,39 +52,10 @@ class NotificationApi(Resource): }, ) @console_ns.response(200, "Success", console_ns.models[NotificationResponse.__name__]) - @setup_required - @login_required - @with_current_user - @account_initialization_required - @only_edition_cloud - def get(self, current_user: Account): - result = BillingService.get_account_notification(str(current_user.id)) - - # Proto JSON uses camelCase field names (Kratos default marshaling). - response: NotificationResponseDict - if not result.get("shouldShow"): - response = {"should_show": False, "notifications": []} - return response, 200 - - lang = current_user.interface_language or _FALLBACK_LANG - - notifications: list[NotificationItemDict] = [] - for notification in result.get("notifications") or []: - contents: Mapping[str, NotificationLangContent] = notification.get("contents") or {} - lang_content = _pick_lang_content(contents, lang) - item: NotificationItemDict = { - "notification_id": notification.get("notificationId"), - "frequency": notification.get("frequency"), - "lang": lang_content.get("lang", lang), - "title": lang_content.get("title", ""), - "subtitle": lang_content.get("subtitle", ""), - "body": lang_content.get("body", ""), - "title_pic_url": lang_content.get("titlePicUrl", ""), - } - notifications.append(item) - - response = {"should_show": bool(notifications), "notifications": notifications} - return response, 200 + @console_account_admission(editions=frozenset({DeploymentEdition.CLOUD})) + def get(self, request_context: RequestContext): + result = application_services().notifications.get_active(request_context) + return dump_response(NotificationResponse, result), 200 @console_ns.route("/notification/dismiss") @@ -134,17 +65,10 @@ class NotificationDismissApi(Resource): description="Mark a notification as dismissed for the current user.", responses={200: "Success", 401: "Unauthorized"}, ) - @setup_required - @login_required - @with_current_user - @account_initialization_required - @only_edition_cloud + @console_account_admission(editions=frozenset({DeploymentEdition.CLOUD})) @console_ns.expect(console_ns.models[DismissNotificationPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) @model_validate(DismissNotificationPayload) - def post(self, payload: DismissNotificationPayload, current_user: Account): - BillingService.dismiss_notification( - notification_id=payload.notification_id, - account_id=str(current_user.id), - ) - return {"result": "success"}, 200 + def post(self, payload: DismissNotificationPayload, request_context: RequestContext): + application_services().notifications.dismiss(request_context, payload.notification_id) + return dump_response(SimpleResultResponse, {"result": "success"}), 200 diff --git a/api/controllers/console/onboarding.py b/api/controllers/console/onboarding.py index f26e2d539e4..cbd77752e7b 100644 --- a/api/controllers/console/onboarding.py +++ b/api/controllers/console/onboarding.py @@ -7,36 +7,20 @@ action-based so callers do not replace server-side arrays with stale snapshots. """ from datetime import datetime -from typing import Literal, cast from flask_restx import Resource from pydantic import BaseModel, ConfigDict, Field, model_validator from controllers.common.schema import register_response_schema_models, register_schema_models -from extensions.ext_database import db +from controllers.console.flask_admission import console_account_admission +from controllers.console.wraps import model_validate +from extensions.ext_application_services import application_services from fields.base import ResponseModel from libs.helper import dump_response -from libs.login import login_required -from models import Account -from services.step_by_step_tour_service import StepByStepTourPatch, StepByStepTourService +from machinery.context import RequestContext +from services.entities.onboarding_entities import StepByStepTourAction, StepByStepTourPatch, StepByStepTourTaskId from . import console_ns -from .wraps import ( - account_initialization_required, - model_validate, - setup_required, - with_current_tenant_id, - with_current_user, -) - -StepByStepTourAction = Literal[ - "skip", - "complete_task", - "uncomplete_task", - "enable_current_workspace", - "disable_current_workspace", -] -StepByStepTourTaskId = Literal["home", "studio", "knowledge", "integration"] class StepByStepTourStatePatchPayload(BaseModel): @@ -74,39 +58,22 @@ class StepByStepTourStateApi(Resource): @console_ns.doc("get_step_by_step_tour_state") @console_ns.doc(description="Get account-level Step-by-step Tour state") @console_ns.response(200, "Success", console_ns.models[StepByStepTourStateResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - def get(self, current_tenant_id: str, current_user: Account): + @console_account_admission() + def get(self, request_context: RequestContext): return dump_response( StepByStepTourStateResponse, - StepByStepTourService.get_state( - account=current_user, - current_tenant_id=current_tenant_id, - session=db.session, - ), + application_services().step_by_step_tour.get_state(request_context), ) @console_ns.doc("patch_step_by_step_tour_state") @console_ns.doc(description="Update account-level Step-by-step Tour state") @console_ns.expect(console_ns.models[StepByStepTourStatePatchPayload.__name__]) @console_ns.response(200, "Success", console_ns.models[StepByStepTourStateResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id + @console_account_admission() @model_validate(StepByStepTourStatePatchPayload) - def patch(self, req_data: StepByStepTourStatePatchPayload, current_tenant_id: str, current_user: Account): - patch = cast(StepByStepTourPatch, req_data.model_dump(exclude_unset=True, exclude_none=True)) + def patch(self, req_data: StepByStepTourStatePatchPayload, request_context: RequestContext): + patch = StepByStepTourPatch(action=req_data.action, task_id=req_data.task_id) return dump_response( StepByStepTourStateResponse, - StepByStepTourService.patch_state( - account=current_user, - current_tenant_id=current_tenant_id, - patch=patch, - session=db.session, - ), + application_services().step_by_step_tour.patch_state(request_context, patch), ) diff --git a/api/dev/generate_swagger_markdown_docs.py b/api/dev/generate_swagger_markdown_docs.py index 991a487c107..a9451c52778 100644 --- a/api/dev/generate_swagger_markdown_docs.py +++ b/api/dev/generate_swagger_markdown_docs.py @@ -76,6 +76,10 @@ def _schema_markdown_type(schema: object) -> str: item_type = _schema_markdown_type(schema.get("items")) return f"[ {item_type or 'object'} ]" if isinstance(schema_type, str): + enum_values = schema.get("enum") + if isinstance(enum_values, list) and enum_values: + rendered_values = ", ".join(json.dumps(value, ensure_ascii=False) for value in enum_values) + return f"{schema_type},
**Available values:** {rendered_values}" return schema_type return "" diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 7a482248cce..747c658a561 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -31,6 +31,7 @@ from repositories.factory import DifyAPIRepositoryFactory from repositories.installation_state_repository import InstallationStateRepository from repositories.oauth_server_repository import RedisOAuthServerTokenRepository, SQLAlchemyOAuthServerRepository from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository +from repositories.step_by_step_tour_repository import SQLAlchemyStepByStepTourStateRepository from repositories.tag_repository import TagRepository from repositories.trial_app_query_repository import TrialAppQueryRepository from repositories.trial_app_usage_repository import TrialAppUsageRepository @@ -104,6 +105,8 @@ 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.notification_gateway import BillingNotificationGateway +from services.notification_service import NotificationService from services.notion_data_source_gateway import NotionDataSourceGateway from services.oauth_server_service import OAUTH_ACCESS_TOKEN_EXPIRES_IN, OAuthServerService from services.partner_tenant_binding_service import PartnerTenantBindingService @@ -122,6 +125,7 @@ from services.retention.workflow_run.archive_log_service import WorkflowRunArchi from services.schema_definition_service import SchemaDefinitionService from services.setup_adapters import RedisSetupLock, RegisterServiceAccountProvisioner from services.setup_service import SetupService +from services.step_by_step_tour_service import StepByStepTourService from services.tag_application_service import TagApplicationService from services.trial_app_usage import TrialAppUsageRecorder from services.web_app_runtime_query_service import WebAppRuntimeQueryService @@ -188,6 +192,8 @@ class ApplicationServices: feature_queries: FeatureQueryService oauth_server: OAuthServerService init_validation: InitValidationService + notifications: NotificationService + step_by_step_tour: StepByStepTourService partner_tenant_bindings: PartnerTenantBindingService recommended_app_queries: RecommendedAppQueryService trial_app_usage: TrialAppUsageRecorder @@ -434,6 +440,16 @@ def build_application_services( validation_required=(deployment_edition != DeploymentEdition.CLOUD and bool(initialization_password)), expected_password=initialization_password, ), + notifications=NotificationService( + accounts=accounts, + notifications=BillingNotificationGateway(), + ), + step_by_step_tour=StepByStepTourService( + accounts=accounts, + states=SQLAlchemyStepByStepTourStateRepository(session_factory=database_client), + enabled=dify_config.ENABLE_STEP_BY_STEP_TOUR, + rollout_started_at=dify_config.STEP_BY_STEP_TOUR_ROLLOUT_STARTED_AT, + ), partner_tenant_bindings=PartnerTenantBindingService( sync_bindings=BillingService.sync_partner_tenants_bindings, ), diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index b8008497e43..c143fd19243 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -13501,7 +13501,7 @@ default (the config form sends the full desired feature state on save). | mode | string,
**Available values:** "advanced-chat", "agent", "agent-chat", "all", "channel", "chat", "completion", "workflow",
**Default:** all | App mode filter
*Enum:* `"advanced-chat"`, `"agent"`, `"agent-chat"`, `"all"`, `"channel"`, `"chat"`, `"completion"`, `"workflow"` | No | | name | string | Filter by app name | No | | page | integer,
**Default:** 1 | Page number (1-99999) | No | -| publication_status | string | Filter by published or draft Agent configuration status | No | +| publication_status | string,
**Available values:** "drafts", "published" | Filter by published or draft Agent configuration status | No | | sort_by | string,
**Available values:** "earliest_created", "last_modified", "recently_created",
**Default:** last_modified | Sort apps by last modified, recently created, or earliest created
*Enum:* `"earliest_created"`, `"last_modified"`, `"recently_created"` | No | | tag_ids | [ string ] | Filter by tag IDs | No | @@ -15744,7 +15744,7 @@ AppMCPServer Status Enum | copyright | string | | No | | custom_disclaimer | string | | No | | customize_domain | string | | No | -| customize_token_strategy | string | | No | +| customize_token_strategy | string,
**Available values:** "allow", "must", "not_allow" | | No | | default_language | string | | No | | description | string | | No | | icon | string | | No | @@ -16202,7 +16202,7 @@ TEAM: Team collaboration paid plan | files | [ object ] | | No | | inputs | object | | Yes | | query | string | | No | -| response_mode | string | | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | | No | | retriever_from | string,
**Default:** explore_app | | No | #### CompletionMessagePayload @@ -16223,7 +16223,7 @@ TEAM: Team collaboration paid plan | files | [ object ] | | No | | inputs | object | | Yes | | query | string | | No | -| response_mode | string | | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | | No | | retriever_from | string,
**Default:** explore_app | | No | #### ComplianceDownloadQuery @@ -18263,9 +18263,9 @@ Flask blueprint initialization. | ---- | ---- | ----------- | -------- | | end_date | string | End date (YYYY-MM-DD) | No | | format | string,
**Available values:** "csv", "json",
**Default:** csv | Export format
*Enum:* `"csv"`, `"json"` | No | -| from_source | string | Filter by feedback source | No | +| from_source | string,
**Available values:** "admin", "user" | Filter by feedback source | No | | has_comment | boolean | Only include feedback with comments | No | -| rating | string | Filter by rating | No | +| rating | string,
**Available values:** "dislike", "like" | Filter by rating | No | | start_date | string | Start date (YYYY-MM-DD) | No | #### FeedbackStat @@ -18663,7 +18663,7 @@ Icon information model. | ---- | ---- | ----------- | -------- | | icon | string | | No | | icon_background | string | | No | -| icon_type | string | | No | +| icon_type | string,
**Available values:** "emoji", "image" | | No | | icon_url | string | | No | #### IconType @@ -19245,7 +19245,7 @@ Enum class for large language model mode. | ---- | ---- | ----------- | -------- | | content | string | Optional text feedback providing additional detail. | No | | message_id | string | Message ID | Yes | -| rating | string | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | +| rating | string,
**Available values:** "dislike", "like" | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | #### MessageFile @@ -19306,7 +19306,7 @@ Metadata Filtering Condition. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | conditions | [ [Condition](#condition) ] | List of metadata conditions to evaluate. | No | -| logical_operator | string | How to combine multiple conditions. | No | +| logical_operator | string,
**Available values:** "and", "or" | How to combine multiple conditions. | No | #### MetadataOperationData @@ -19439,7 +19439,7 @@ Enum class for model property key. | is_exhausted | boolean | | Yes | | is_unlimited | boolean | | Yes | | next_credit_reset_date | integer | | Yes | -| pool_type | string | | Yes | +| pool_type | string,
**Available values:** "paid", "trial" | | Yes | | quota_limit | integer | Credit limit for the effective pool; -1 means unlimited. | Yes | | quota_used | integer | | Yes | | remaining_credits | integer | Remaining credits; -1 means unlimited. | Yes | @@ -21434,7 +21434,7 @@ Model class for provider quota configuration. | ---- | ---- | ----------- | -------- | | metadata_filtering_conditions | [MetadataFilteringCondition](#metadatafilteringcondition) | Restrict retrieval to chunks whose document metadata matches the given conditions. Conditions are evaluated server-side against document metadata fields. | No | | reranking_enable | boolean | Whether reranking is enabled. | Yes | -| reranking_mode | string | Reranking mode. Required when `reranking_enable` is `true`. | No | +| reranking_mode | string,
**Available values:** "reranking_model", "weighted_score" | Reranking mode. Required when `reranking_enable` is `true`. | No | | reranking_model | [RerankingModel](#rerankingmodel) | Reranking model configuration. | No | | score_threshold | number | Minimum similarity score for results. Only effective when score threshold filtering is enabled. | No | | score_threshold_enabled | boolean | Whether score threshold filtering is enabled. | Yes | @@ -21488,7 +21488,7 @@ Model class for provider quota configuration. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| parent_mode | string | Parent-child segmentation mode. | No | +| parent_mode | string,
**Available values:** "full-doc", "paragraph" | Parent-child segmentation mode. | No | | pre_processing_rules | [ [PreProcessingRule](#preprocessingrule) ] | Pre-processing rules to apply before segmentation. | No | | segmentation | [Segmentation](#segmentation) | Parent chunk segmentation settings. | No | | subchunk_segmentation | [Segmentation](#segmentation) | Child chunk segmentation settings. | No | @@ -22477,7 +22477,7 @@ Query parameters for listing snippet published workflows. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | action | string,
**Available values:** "complete_task", "disable_current_workspace", "enable_current_workspace", "skip", "uncomplete_task" | State update action
*Enum:* `"complete_task"`, `"disable_current_workspace"`, `"enable_current_workspace"`, `"skip"`, `"uncomplete_task"` | Yes | -| task_id | string | Task ID for task actions | No | +| task_id | string,
**Available values:** "home", "integration", "knowledge", "studio" | Task ID for task actions | No | #### StepByStepTourStateResponse @@ -22943,7 +22943,7 @@ Tool label | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| visibility | string | Visibility for the OAuth credential. Defaults to 'only_me'. | No | +| visibility | string,
**Available values:** "all_team_members", "only_me" | Visibility for the OAuth credential. Defaults to 'only_me'. | No | #### ToolOAuthCustomClientPayload @@ -23075,7 +23075,7 @@ removes TOOLS_SELECTOR from PluginParameterType | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| type | string | | No | +| type | string,
**Available values:** "api", "builtin", "mcp", "model", "workflow" | | No | #### ToolProviderListResponse @@ -23693,7 +23693,7 @@ in form definition, or a variable while the workflow is running. | ---- | ---- | ----------- | -------- | | keyword_setting | [WeightKeywordSetting](#weightkeywordsetting) | Keyword search weight settings. | No | | vector_setting | [WeightVectorSetting](#weightvectorsetting) | Semantic search weight settings. | No | -| weight_type | string | Strategy for balancing semantic and keyword search weights. | No | +| weight_type | string,
**Available values:** "customized", "keyword_first", "semantic_first" | Strategy for balancing semantic and keyword search weights. | No | #### WeightVectorSetting @@ -24199,7 +24199,7 @@ can reuse its existing handler. | description | string | | No | | event | string | | No | | icon | string | | No | -| mode | string | *Enum:* `"advanced-chat"`, `"workflow"` | Yes | +| mode | string,
**Available values:** "advanced-chat", "workflow" | *Enum:* `"advanced-chat"`, `"workflow"` | Yes | | nodes | [ [WorkflowPlanNodeResponse](#workflowplannoderesponse) ] | | Yes | | start_inputs | [ [WorkflowPlanStartInputResponse](#workflowplanstartinputresponse) ] | | No | | title | string | | No | @@ -24214,7 +24214,7 @@ can reuse its existing handler. | graph | [WorkflowGraph](#workflowgraph) | | Yes | | icon | string | | No | | message | string | | No | -| mode | string | | No | +| mode | string,
**Available values:** "advanced-chat", "workflow" | | No | #### WorkflowGenerateResultEventResponse @@ -24227,7 +24227,7 @@ can reuse its existing handler. | graph | [WorkflowGraph](#workflowgraph) | | Yes | | icon | string | | No | | message | string | | No | -| mode | string | | No | +| mode | string,
**Available values:** "advanced-chat", "workflow" | | No | #### WorkflowGenerateStreamEventResponse @@ -24527,9 +24527,9 @@ Lifecycle state for an asynchronous archive download request. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| status | string | Workflow run status filter | No | +| status | string,
**Available values:** "failed", "partial-succeeded", "running", "stopped", "succeeded" | Workflow run status filter | No | | time_range | string | Filter by time range (optional): e.g., 7d (7 days), 4h (4 hours), 30m (30 minutes), 30s (30 seconds). Filters by created_at field. | No | -| triggered_from | string | Filter by trigger source: debugging or app-run. Default: debugging | No | +| triggered_from | string,
**Available values:** "app-run", "debugging" | Filter by trigger source: debugging or app-run. Default: debugging | No | #### WorkflowRunCountResponse @@ -24601,8 +24601,8 @@ Lifecycle state for an asynchronous archive download request. | ---- | ---- | ----------- | -------- | | last_id | string | Last run ID for pagination | No | | limit | integer,
**Default:** 20 | Number of items per page (1-100) | No | -| status | string | Workflow run status filter | No | -| triggered_from | string | Filter by trigger source: debugging or app-run. Default: debugging | No | +| status | string,
**Available values:** "failed", "partial-succeeded", "running", "stopped", "succeeded" | Workflow run status filter | No | +| triggered_from | string,
**Available values:** "app-run", "debugging" | Filter by trigger source: debugging or app-run. Default: debugging | No | #### WorkflowRunNodeExecutionListResponse @@ -24900,7 +24900,7 @@ Workflow tool configuration | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| language | string | Localized policy label language | No | +| language | string,
**Available values:** "en", "ja", "zh" | Localized policy label language | No | #### _AccessPolicyList @@ -24959,7 +24959,7 @@ Workflow tool configuration | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| language | string | Localized policy label language | No | +| language | string,
**Available values:** "en", "ja", "zh" | Localized policy label language | No | | limit | integer | | No | | page | integer | | No | | reverse | boolean | | No | diff --git a/api/openapi/markdown/service-openapi.md b/api/openapi/markdown/service-openapi.md index a09bd57e255..80b5ab94db8 100644 --- a/api/openapi/markdown/service-openapi.md +++ b/api/openapi/markdown/service-openapi.md @@ -2587,7 +2587,7 @@ Public pause reason emitted by a blocking Chatflow execution. | files | [ object
object
object
object ] | File list for multimodal understanding, including images, documents, audio, and video. To attach a local file, first upload it via [Upload File](/api-reference/files/upload-file) and use the returned `id` as `upload_file_id` with `transfer_method: local_file`. | No | | inputs | object | Values for app-defined variables. Refer to the `user_input_form` field in the [Get App Parameters](/api-reference/applications/get-app-parameters) response to discover expected variable names and types. | Yes | | query | string | User input or question content. | Yes | -| response_mode | string | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. New Agent app mode supports streaming only. When omitted, non-Agent apps run in blocking mode and new Agent apps stream. | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. New Agent app mode supports streaming only. When omitted, non-Agent apps run in blocking mode and new Agent apps stream. | No | | workflow_id | string | Published workflow version ID to execute for advanced chat. If omitted, the app's current published workflow is used. | No | #### ChatRequestPayloadWithUser @@ -2599,7 +2599,7 @@ Public pause reason emitted by a blocking Chatflow execution. | files | [ object
object
object
object ] | File list for multimodal understanding, including images, documents, audio, and video. To attach a local file, first upload it via [Upload File](/api-reference/files/upload-file) and use the returned `id` as `upload_file_id` with `transfer_method: local_file`. | No | | inputs | object | Values for app-defined variables. Refer to the `user_input_form` field in the [Get App Parameters](/api-reference/applications/get-app-parameters) response to discover expected variable names and types. | Yes | | query | string | User input or question content. | Yes | -| response_mode | string | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. New Agent app mode supports streaming only. When omitted, non-Agent apps run in blocking mode and new Agent apps stream. | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. New Agent app mode supports streaming only. When omitted, non-Agent apps run in blocking mode and new Agent apps stream. | No | | user | string | User identifier, unique within the application. This identifier scopes data access; resources created with one `user` value are only visible when queried with the same `user` value. | Yes | | workflow_id | string | Published workflow version ID to execute for advanced chat. If omitted, the app's current published workflow is used. | No | @@ -2672,7 +2672,7 @@ Public pause reason emitted by a blocking Chatflow execution. | files | [ object
object
object
object ] | File list for multimodal understanding, including images, documents, audio, and video. To attach a local file, first upload it via [Upload File](/api-reference/files/upload-file) and use the returned `id` as `upload_file_id` with `transfer_method: local_file`. | No | | inputs | object | Values for app-defined variables. Refer to the `user_input_form` field in the [Get App Parameters](/api-reference/applications/get-app-parameters) response to discover expected variable names and types. | Yes | | query | string | User input or prompt content. | No | -| response_mode | string | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. When omitted, the request runs in blocking mode. | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. When omitted, the request runs in blocking mode. | No | #### CompletionRequestPayloadWithUser @@ -2681,7 +2681,7 @@ Public pause reason emitted by a blocking Chatflow execution. | files | [ object
object
object
object ] | File list for multimodal understanding, including images, documents, audio, and video. To attach a local file, first upload it via [Upload File](/api-reference/files/upload-file) and use the returned `id` as `upload_file_id` with `transfer_method: local_file`. | No | | inputs | object | Values for app-defined variables. Refer to the `user_input_form` field in the [Get App Parameters](/api-reference/applications/get-app-parameters) response to discover expected variable names and types. | Yes | | query | string | User input or prompt content. | No | -| response_mode | string | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. When omitted, the request runs in blocking mode. | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode. `streaming` uses Server-Sent Events; `blocking` returns after completion. When omitted, the request runs in blocking mode. | No | | user | string | User identifier, unique within the application. This identifier scopes data access; resources created with one `user` value are only visible when queried with the same `user` value. | Yes | #### Condition @@ -2797,7 +2797,7 @@ Enum class for custom configuration status. | embedding_model_provider | string | Embedding model provider. Use the `provider` field from [Get Available Models](/api-reference/models/get-available-models) with `model_type=text-embedding`. | No | | external_knowledge_api_id | string | ID of the external knowledge API. | No | | external_knowledge_id | string | ID of the external knowledge base. | No | -| indexing_technique | string | `high_quality` uses embedding models for precise search; `economy` uses keyword-based indexing. | No | +| indexing_technique | string,
**Available values:** "economy", "high_quality" | `high_quality` uses embedding models for precise search; `economy` uses keyword-based indexing. | No | | name | string | Name of the knowledge base. | Yes | | permission | [PermissionEnum](#permissionenum) | Controls who can access this knowledge base. `only_me` restricts access to the creator, `all_team_members` grants workspace-wide access, and `partial_members` grants access to specified members. | No | | provider | string,
**Available values:** "external", "vendor",
**Default:** vendor | Knowledge base provider: `vendor` for internal knowledge bases, `external` for external ones.
*Enum:* `"external"`, `"vendor"` | No | @@ -3039,7 +3039,7 @@ Enum class for custom configuration status. | external_knowledge_api_id | string | ID of the external knowledge API. | No | | external_knowledge_id | string | ID of the external knowledge base. | No | | external_retrieval_model | object | Retrieval settings for external knowledge bases. | No | -| indexing_technique | string | `high_quality` uses embedding models for precise search; `economy` uses keyword-based indexing. | No | +| indexing_technique | string,
**Available values:** "economy", "high_quality" | `high_quality` uses embedding models for precise search; `economy` uses keyword-based indexing. | No | | name | string | Name of the knowledge base. | No | | partial_member_list | [ object ] | List of team members with access when `permission` is `partial_members`. | No | | permission | [PermissionEnum](#permissionenum) | Controls who can access this knowledge base. `only_me` restricts access to the creator, `all_team_members` grants workspace-wide access, and `partial_members` grants access to specified members. | No | @@ -3167,7 +3167,7 @@ Request payload for bulk downloading documents as a zip archive. | keyword | string | Search keyword to filter by document name. | No | | limit | integer,
**Default:** 20 | Number of items per page. Server caps at `100`. | No | | page | integer,
**Default:** 1 | Page number to retrieve. | No | -| status | string | Filter by display status. | No | +| status | string,
**Available values:** "archived", "available", "disabled", "error", "indexing", "paused", "queuing" | Filter by display status. | No | #### DocumentListResponse @@ -3265,7 +3265,7 @@ Request payload for bulk downloading documents as a zip archive. | doc_language | string,
**Default:** English | Language of the document for processing optimization. | No | | embedding_model | string | Embedding model name. Use the `model` field from [Get Available Models](/api-reference/models/get-available-models) with `model_type=text-embedding`. | No | | embedding_model_provider | string | Embedding model provider. Use the `provider` field from [Get Available Models](/api-reference/models/get-available-models) with `model_type=text-embedding`. | No | -| indexing_technique | string | `high_quality` uses embedding models for precise search; `economy` uses keyword-based indexing. Required when adding the first document to a knowledge base; subsequent documents inherit the knowledge base's indexing technique if omitted. | No | +| indexing_technique | string,
**Available values:** "economy", "high_quality" | `high_quality` uses embedding models for precise search; `economy` uses keyword-based indexing. Required when adding the first document to a knowledge base; subsequent documents inherit the knowledge base's indexing technique if omitted. | No | | name | string | Document name. | Yes | | original_document_id | string | Original document ID for replacement. | No | | process_rule | [ProcessRule](#processrule) | Processing rules for chunking. | No | @@ -3614,14 +3614,14 @@ Model class for i18n object. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | content | string | Optional text feedback providing additional detail. | No | -| rating | string | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | +| rating | string,
**Available values:** "dislike", "like" | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | #### MessageFeedbackPayloadWithUser | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | content | string | Optional text feedback providing additional detail. | No | -| rating | string | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | +| rating | string,
**Available values:** "dislike", "like" | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | | user | string | User identifier, unique within the application. This identifier scopes data access; resources created with one `user` value are only visible when queried with the same `user` value. | Yes | #### MessageFile @@ -3701,7 +3701,7 @@ Metadata Filtering Condition. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | conditions | [ [Condition](#condition) ] | List of metadata conditions to evaluate. | No | -| logical_operator | string | How to combine multiple conditions. | No | +| logical_operator | string,
**Available values:** "and", "or" | How to combine multiple conditions. | No | #### MetadataOperationData @@ -3935,7 +3935,7 @@ Model class for provider with models response. | ---- | ---- | ----------- | -------- | | metadata_filtering_conditions | [MetadataFilteringCondition](#metadatafilteringcondition) | Restrict retrieval to chunks whose document metadata matches the given conditions. Conditions are evaluated server-side against document metadata fields. | No | | reranking_enable | boolean | Whether reranking is enabled. | Yes | -| reranking_mode | string | Reranking mode. Required when `reranking_enable` is `true`. | No | +| reranking_mode | string,
**Available values:** "reranking_model", "weighted_score" | Reranking mode. Required when `reranking_enable` is `true`. | No | | reranking_model | [RerankingModel](#rerankingmodel) | Reranking model configuration. | No | | score_threshold | number | Minimum similarity score for results. Only effective when score threshold filtering is enabled. | No | | score_threshold_enabled | boolean | Whether score threshold filtering is enabled. | Yes | @@ -3969,7 +3969,7 @@ Model class for provider with models response. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| parent_mode | string | Parent-child segmentation mode. | No | +| parent_mode | string,
**Available values:** "full-doc", "paragraph" | Parent-child segmentation mode. | No | | pre_processing_rules | [ [PreProcessingRule](#preprocessingrule) ] | Pre-processing rules to apply before segmentation. | No | | segmentation | [Segmentation](#segmentation) | Parent chunk segmentation settings. | No | | subchunk_segmentation | [Segmentation](#segmentation) | Child chunk segmentation settings. | No | @@ -4300,7 +4300,7 @@ in form definition, or a variable while the workflow is running. | ---- | ---- | ----------- | -------- | | keyword_setting | [WeightKeywordSetting](#weightkeywordsetting) | Keyword search weight settings. | No | | vector_setting | [WeightVectorSetting](#weightvectorsetting) | Semantic search weight settings. | No | -| weight_type | string | Strategy for balancing semantic and keyword search weights. | No | +| weight_type | string,
**Available values:** "customized", "keyword_first", "semantic_first" | Strategy for balancing semantic and keyword search weights. | No | #### WeightVectorSetting @@ -4383,7 +4383,7 @@ Blocking workflow response for a finished or paused execution. | keyword | string | Keyword to search in logs. | No | | limit | integer,
**Default:** 20 | Number of items per page. | No | | page | integer,
**Default:** 1 | Page number for pagination. | No | -| status | string | Filter by execution status. | No | +| status | string,
**Available values:** "failed", "stopped", "succeeded" | Filter by execution status. | No | #### WorkflowPauseReasonResponse @@ -4452,7 +4452,7 @@ Public pause reason emitted by a blocking Workflow execution. | ---- | ---- | ----------- | -------- | | files | [ object
object
object
object ] | File list for workflow system file inputs. Available when file upload is enabled for the workflow. To attach a local file, first upload it via [Upload File](/api-reference/files/upload-file) and use the returned `id` as `upload_file_id` with `transfer_method: local_file`. | No | | inputs | object | Key-value pairs for workflow input variables. Values for file-type variables should be arrays of file objects with `type`, `transfer_method`, and either `url` or `upload_file_id`. Refer to the `user_input_form` field in the [Get App Parameters](/api-reference/applications/get-app-parameters) response to discover the variable names and types expected by your app. | Yes | -| response_mode | string | Response mode. Use `blocking` for synchronous responses or `streaming` for Server-Sent Events. When omitted, the request runs in blocking mode. | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode. Use `blocking` for synchronous responses or `streaming` for Server-Sent Events. When omitted, the request runs in blocking mode. | No | #### WorkflowRunPayloadWithUser @@ -4460,7 +4460,7 @@ Public pause reason emitted by a blocking Workflow execution. | ---- | ---- | ----------- | -------- | | files | [ object
object
object
object ] | File list for workflow system file inputs. Available when file upload is enabled for the workflow. To attach a local file, first upload it via [Upload File](/api-reference/files/upload-file) and use the returned `id` as `upload_file_id` with `transfer_method: local_file`. | No | | inputs | object | Key-value pairs for workflow input variables. Values for file-type variables should be arrays of file objects with `type`, `transfer_method`, and either `url` or `upload_file_id`. Refer to the `user_input_form` field in the [Get App Parameters](/api-reference/applications/get-app-parameters) response to discover the variable names and types expected by your app. | Yes | -| response_mode | string | Response mode. Use `blocking` for synchronous responses or `streaming` for Server-Sent Events. When omitted, the request runs in blocking mode. | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode. Use `blocking` for synchronous responses or `streaming` for Server-Sent Events. When omitted, the request runs in blocking mode. | No | | user | string | User identifier, unique within the application. This identifier scopes data access; resources created with one `user` value are only visible when queried with the same `user` value. | Yes | #### WorkflowRunResponse diff --git a/api/openapi/markdown/web-openapi.md b/api/openapi/markdown/web-openapi.md index e534fe39350..3cc99e099ce 100644 --- a/api/openapi/markdown/web-openapi.md +++ b/api/openapi/markdown/web-openapi.md @@ -1019,7 +1019,7 @@ Button styles for user actions. | inputs | object | Input variables for the chat | Yes | | parent_message_id | string | Parent message ID | No | | query | string | User query/message | Yes | -| response_mode | string | Response mode: blocking or streaming | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode: blocking or streaming | No | | retriever_from | string,
**Default:** web_app | Source of retriever | No | #### CompletionMessagePayload @@ -1029,7 +1029,7 @@ Button styles for user actions. | files | [ object ] | Files to be processed | No | | inputs | object | Input variables for the completion | Yes | | query | string | Query text for completion | No | -| response_mode | string | Response mode: blocking or streaming | No | +| response_mode | string,
**Available values:** "blocking", "streaming" | Response mode: blocking or streaming | No | | retriever_from | string,
**Default:** web_app | Source of retriever | No | #### ConversationInfiniteScrollPagination @@ -1322,7 +1322,7 @@ Parsed multipart form fields for HITL uploads. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | content | string | Optional text feedback providing additional detail. | No | -| rating | string | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | +| rating | string,
**Available values:** "dislike", "like" | Feedback rating. Set to `null` to revoke previously submitted feedback. | No | #### MessageFile diff --git a/api/repositories/step_by_step_tour_repository.py b/api/repositories/step_by_step_tour_repository.py new file mode 100644 index 00000000000..7dc7a6d6bf2 --- /dev/null +++ b/api/repositories/step_by_step_tour_repository.py @@ -0,0 +1,189 @@ +"""SQLAlchemy repository for account Step-by-step Tour state.""" + +import logging +from collections.abc import Callable +from typing import Protocol, override, runtime_checkable + +from sqlalchemy import select, update +from sqlalchemy.exc import IntegrityError, OperationalError +from sqlalchemy.orm import Session, sessionmaker + +from models.onboarding import AccountStepByStepTourState +from services.entities.onboarding_entities import StepByStepTourState +from services.step_by_step_tour_service import StepByStepTourStateRepository + +logger = logging.getLogger(__name__) + +_MYSQL_RETRYABLE_LOCK_ERRNOS = frozenset({1205, 1213}) +_MAX_LOCK_ATTEMPTS = 3 + + +@runtime_checkable +class _ErrorWithErrno(Protocol): + @property + def errno(self) -> object: ... + + +class SQLAlchemyStepByStepTourStateRepository(StepByStepTourStateRepository): + def __init__(self, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def get(self, account_id: str) -> StepByStepTourState | None: + with self._session_factory() as session: + model = self._get_model(account_id, session=session) + return self._to_state(model) if model is not None else None + + @override + def initialize(self, account_id: str, first_workspace_id: str) -> StepByStepTourState: + """Create state with its first workspace, or atomically claim a legacy empty state.""" + return self._run_with_lock_retry( + lambda: self._initialize_once(account_id, first_workspace_id), + ) + + def _initialize_once(self, account_id: str, first_workspace_id: str) -> StepByStepTourState: + with self._session_factory() as session: + model = self._get_model(account_id, session=session) + if model is None: + model = AccountStepByStepTourState( + account_id=account_id, + first_workspace_id=first_workspace_id, + ) + session.add(model) + try: + session.commit() + except IntegrityError: + # A concurrent request inserted the account-owned row first. + session.rollback() + model = self._get_model(account_id, session=session) + if model is None: + raise + else: + session.refresh(model) + return self._to_state(model) + + if model.first_workspace_id is None: + stmt = ( + update(AccountStepByStepTourState) + .where( + AccountStepByStepTourState.account_id == account_id, + AccountStepByStepTourState.first_workspace_id.is_(None), + ) + .values(first_workspace_id=first_workspace_id) + .execution_options(synchronize_session=False) + ) + session.execute(stmt) + session.commit() + # A competing conditional update may have won while this request waited. + session.refresh(model) + + return self._to_state(model) + + @override + def mutate( + self, + account_id: str, + mutation: Callable[[StepByStepTourState], StepByStepTourState], + ) -> StepByStepTourState: + """Lock, create if needed, mutate, and persist account state in one transaction.""" + return self._run_with_lock_retry( + lambda: self._mutate_once(account_id, mutation), + ) + + def _mutate_once( + self, + account_id: str, + mutation: Callable[[StepByStepTourState], StepByStepTourState], + ) -> StepByStepTourState: + with self._session_factory() as session: + # Probe without a locking read so a missing MySQL unique key does not + # acquire a gap/next-key lock before the insert. + model = self._get_model(account_id, session=session) + if model is None: + model = AccountStepByStepTourState(account_id=account_id) + session.add(model) + try: + session.flush() + except IntegrityError: + # A concurrent mutation created the row. Start a new transaction, + # lock its committed state, and replay the pure mutation on it. + session.rollback() + model = self._get_model(account_id, session=session, lock_for_update=True) + if model is None: + raise + else: + model = self._get_model(account_id, session=session, lock_for_update=True) + if model is None: + raise RuntimeError("Step-by-step Tour state disappeared while acquiring its lock") + + state = mutation(self._to_state(model)) + if state.account_id != account_id: + raise ValueError("Step-by-step Tour mutation cannot change account ownership") + # first_workspace_id is write-once and owned exclusively by initialize(). + model.skipped = state.skipped + model.completed_task_ids = list(state.completed_task_ids) + model.manually_enabled_workspace_ids = list(state.manually_enabled_workspace_ids) + model.manually_disabled_workspace_ids = list(state.manually_disabled_workspace_ids) + session.commit() + session.refresh(model) + return self._to_state(model) + + @staticmethod + def _run_with_lock_retry[T](operation: Callable[[], T]) -> T: + for attempt in range(1, _MAX_LOCK_ATTEMPTS): + try: + return operation() + except OperationalError as exc: + if not _is_retryable_mysql_lock_error(exc): + raise + logger.warning( + "Retrying Step-by-step Tour transaction after MySQL lock failure (attempt %s/%s)", + attempt, + _MAX_LOCK_ATTEMPTS, + ) + return operation() + + @staticmethod + def _get_model( + account_id: str, + *, + session: Session, + lock_for_update: bool = False, + ) -> AccountStepByStepTourState | None: + stmt = select(AccountStepByStepTourState).where(AccountStepByStepTourState.account_id == account_id).limit(1) + if lock_for_update: + stmt = stmt.with_for_update().execution_options(populate_existing=True) + return session.execute(stmt).scalar_one_or_none() + + @staticmethod + def _to_state(model: AccountStepByStepTourState) -> StepByStepTourState: + return StepByStepTourState( + account_id=model.account_id, + first_workspace_id=model.first_workspace_id, + skipped=model.skipped, + completed_task_ids=tuple(model.completed_task_ids), + manually_enabled_workspace_ids=tuple(model.manually_enabled_workspace_ids), + manually_disabled_workspace_ids=tuple(model.manually_disabled_workspace_ids), + updated_at=model.updated_at, + ) + + +def _is_retryable_mysql_lock_error(exc: OperationalError) -> bool: + orig = exc.orig + if isinstance(orig, _ErrorWithErrno) and _is_retryable_mysql_lock_error_code(orig.errno): + return True + if not isinstance(orig, BaseException) or not orig.args: + return False + return _is_retryable_mysql_lock_error_code(orig.args[0]) + + +def _is_retryable_mysql_lock_error_code(candidate: object) -> bool: + if isinstance(candidate, bool): + return False + if isinstance(candidate, int): + code = candidate + elif isinstance(candidate, str) and candidate.isdecimal(): + code = int(candidate) + else: + return False + return code in _MYSQL_RETRYABLE_LOCK_ERRNOS diff --git a/api/services/entities/notification_entities.py b/api/services/entities/notification_entities.py new file mode 100644 index 00000000000..6686c5edb99 --- /dev/null +++ b/api/services/entities/notification_entities.py @@ -0,0 +1,38 @@ +"""Framework-independent notification contracts.""" + +from collections.abc import Mapping +from typing import NamedTuple + + +class NotificationContent(NamedTuple): + lang: str + title: str + subtitle: str + body: str + title_pic_url: str + + +class AccountNotification(NamedTuple): + notification_id: str | None + frequency: str | None + contents: Mapping[str, NotificationContent] + + +class AccountNotificationBatch(NamedTuple): + should_show: bool + notifications: tuple[AccountNotification, ...] + + +class NotificationItem(NamedTuple): + notification_id: str | None + frequency: str | None + lang: str + title: str + subtitle: str + body: str + title_pic_url: str + + +class NotificationResult(NamedTuple): + should_show: bool + notifications: tuple[NotificationItem, ...] diff --git a/api/services/entities/onboarding_entities.py b/api/services/entities/onboarding_entities.py new file mode 100644 index 00000000000..2550db489e4 --- /dev/null +++ b/api/services/entities/onboarding_entities.py @@ -0,0 +1,42 @@ +"""Framework-independent Step-by-step Tour contracts.""" + +from dataclasses import dataclass +from datetime import datetime +from typing import Literal, TypeAlias + +# Assignment-form aliases preserve Literal enum values in Pydantic-generated OpenAPI schemas. +StepByStepTourAction: TypeAlias = Literal[ # noqa: UP040 + "skip", + "complete_task", + "uncomplete_task", + "enable_current_workspace", + "disable_current_workspace", +] +StepByStepTourTaskId: TypeAlias = Literal["home", "studio", "knowledge", "integration"] # noqa: UP040 + + +@dataclass(frozen=True, slots=True) +class StepByStepTourPatch: + action: StepByStepTourAction + task_id: StepByStepTourTaskId | None = None + + +@dataclass(frozen=True, slots=True) +class StepByStepTourState: + account_id: str + first_workspace_id: str | None = None + skipped: bool = False + completed_task_ids: tuple[str, ...] = () + manually_enabled_workspace_ids: tuple[str, ...] = () + manually_disabled_workspace_ids: tuple[str, ...] = () + updated_at: datetime | None = None + + +@dataclass(frozen=True, slots=True) +class StepByStepTourResult: + first_workspace_id: str | None = None + skipped: bool = False + completed_task_ids: tuple[str, ...] = () + manually_enabled_workspace_ids: tuple[str, ...] = () + manually_disabled_workspace_ids: tuple[str, ...] = () + updated_at: datetime | None = None diff --git a/api/services/notification_gateway.py b/api/services/notification_gateway.py new file mode 100644 index 00000000000..cb7cc5e74d0 --- /dev/null +++ b/api/services/notification_gateway.py @@ -0,0 +1,48 @@ +"""Billing-backed notification gateway.""" + +from collections.abc import Mapping +from typing import Any, override + +from services.billing_service import BillingService +from services.entities.notification_entities import ( + AccountNotification, + AccountNotificationBatch, + NotificationContent, +) +from services.notification_service import NotificationGateway + + +class BillingNotificationGateway(NotificationGateway): + @override + def get_active(self, account_id: str) -> AccountNotificationBatch: + payload = BillingService.get_account_notification(account_id) + notifications = tuple(self._map_notification(item) for item in payload.get("notifications") or ()) + return AccountNotificationBatch( + should_show=bool(payload.get("shouldShow")), + notifications=notifications, + ) + + @override + def dismiss(self, notification_id: str, account_id: str) -> None: + BillingService.dismiss_notification(notification_id=notification_id, account_id=account_id) + + @classmethod + def _map_notification(cls, payload: Mapping[str, Any]) -> AccountNotification: + raw_contents = payload.get("contents") or {} + contents = {language: cls._map_content(content) for language, content in raw_contents.items() if content} + return AccountNotification( + notification_id=payload.get("notificationId"), + frequency=payload.get("frequency"), + contents=contents, + ) + + @staticmethod + def _map_content(payload: Mapping[str, Any]) -> NotificationContent: + return NotificationContent( + # The application service owns the requested-language fallback. + lang=payload.get("lang") or "", + title=payload.get("title") or "", + subtitle=payload.get("subtitle") or "", + body=payload.get("body") or "", + title_pic_url=payload.get("titlePicUrl") or "", + ) diff --git a/api/services/notification_service.py b/api/services/notification_service.py new file mode 100644 index 00000000000..13236ef16ed --- /dev/null +++ b/api/services/notification_service.py @@ -0,0 +1,60 @@ +"""Application service for Console account notifications.""" + +from typing import Protocol + +from machinery.context import RequestContext +from services.account_ports import AccountRepository +from services.entities.notification_entities import ( + AccountNotification, + AccountNotificationBatch, + NotificationContent, + NotificationItem, + NotificationResult, +) + +_FALLBACK_LANGUAGE = "en-US" + + +class NotificationGateway(Protocol): + def get_active(self, account_id: str) -> AccountNotificationBatch: ... + + def dismiss(self, notification_id: str, account_id: str) -> None: ... + + +class NotificationService: + def __init__(self, *, accounts: AccountRepository, notifications: NotificationGateway) -> None: + self._accounts = accounts + self._notifications = notifications + + def get_active(self, context: RequestContext) -> NotificationResult: + batch = self._notifications.get_active(context.account_id) + if not batch.should_show: + return NotificationResult(should_show=False, notifications=()) + + account = self._accounts.get(context.account_id) + if account is None: + raise RuntimeError("Console account admission resolved an unknown account") + language = account.interface_language or _FALLBACK_LANGUAGE + + notifications = tuple(self._localize(notification, language) for notification in batch.notifications) + return NotificationResult(should_show=bool(notifications), notifications=notifications) + + def dismiss(self, context: RequestContext, notification_id: str) -> None: + self._notifications.dismiss(notification_id, context.account_id) + + @staticmethod + def _localize(notification: AccountNotification, language: str) -> NotificationItem: + content = ( + notification.contents.get(language) + or notification.contents.get(_FALLBACK_LANGUAGE) + or next(iter(notification.contents.values()), NotificationContent(language, "", "", "", "")) + ) + return NotificationItem( + notification_id=notification.notification_id, + frequency=notification.frequency, + lang=content.lang or language, + title=content.title, + subtitle=content.subtitle, + body=content.body, + title_pic_url=content.title_pic_url, + ) diff --git a/api/services/step_by_step_tour_service.py b/api/services/step_by_step_tour_service.py index b01d59c1acc..9597d3d5e77 100644 --- a/api/services/step_by_step_tour_service.py +++ b/api/services/step_by_step_tour_service.py @@ -1,221 +1,161 @@ -"""Account-level Step-by-step Tour persistence.""" +"""Application service for account-level Step-by-step Tour use cases.""" +from collections.abc import Callable +from dataclasses import replace from datetime import datetime -from typing import NotRequired, TypedDict +from typing import Protocol, get_args -from sqlalchemy import select -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session, scoped_session - -from configs import dify_config from libs.datetime_utils import ensure_naive_utc -from models.account import Account -from models.onboarding import AccountStepByStepTourState +from machinery.context import RequestContext +from services.account_ports import AccountRepository +from services.entities.onboarding_entities import ( + StepByStepTourPatch, + StepByStepTourResult, + StepByStepTourState, + StepByStepTourTaskId, +) -STEP_BY_STEP_TOUR_TASK_IDS = frozenset(("home", "studio", "knowledge", "integration")) +_TASK_IDS: frozenset[str] = frozenset(get_args(StepByStepTourTaskId)) -class StepByStepTourStateResponse(TypedDict): - first_workspace_id: str | None - skipped: bool - completed_task_ids: list[str] - manually_enabled_workspace_ids: list[str] - manually_disabled_workspace_ids: list[str] - updated_at: datetime | None +class StepByStepTourStateRepository(Protocol): + def get(self, account_id: str) -> StepByStepTourState | None: ... + def initialize(self, account_id: str, first_workspace_id: str) -> StepByStepTourState: ... -class StepByStepTourPatch(TypedDict): - action: str - task_id: NotRequired[str | None] + def mutate( + self, + account_id: str, + mutation: Callable[[StepByStepTourState], StepByStepTourState], + ) -> StepByStepTourState: ... class StepByStepTourService: - """Coordinate persisted tour state with account eligibility rules.""" - - @classmethod - def get_state( - cls, + def __init__( + self, *, - account: Account, - current_tenant_id: str, - session: Session | scoped_session, - ) -> StepByStepTourStateResponse: - eligible = cls.is_eligible(account) - state = cls._get_state(account.id, session=session) + accounts: AccountRepository, + states: StepByStepTourStateRepository, + enabled: bool, + rollout_started_at: datetime | None, + ) -> None: + self._accounts = accounts + self._states = states + self._enabled = enabled + self._rollout_started_at = rollout_started_at - if eligible: - state = cls._ensure_state(account.id, session=session, state=state) - if state.first_workspace_id is None: - state.first_workspace_id = current_tenant_id - session.commit() - session.refresh(state) + def get_state(self, context: RequestContext) -> StepByStepTourResult: + workspace_id = self._require_workspace(context) + account = self._accounts.get(context.account_id) + if account is None: + raise RuntimeError("Console account admission resolved an unknown account") - return cls._build_response(state=state) + if not self._is_eligible(account.initialized_at or account.created_at): + return self._to_result(self._states.get(context.account_id)) - @classmethod - def patch_state( - cls, - *, - account: Account, - current_tenant_id: str, - patch: StepByStepTourPatch, - session: Session | scoped_session, - ) -> StepByStepTourStateResponse: - state = cls._ensure_state(account.id, session=session, state=None) - cls._apply_action( - state=state, - action=patch["action"], - task_id=patch.get("task_id"), - current_tenant_id=current_tenant_id, + return self._to_result(self._states.initialize(context.account_id, workspace_id)) + + def patch_state(self, context: RequestContext, patch: StepByStepTourPatch) -> StepByStepTourResult: + workspace_id = self._require_workspace(context) + state = self._states.mutate( + context.account_id, + lambda current: self._apply_action(current, patch=patch, workspace_id=workspace_id), ) + return self._to_result(state) - session.commit() - session.refresh(state) - return cls._build_response(state=state) - - @classmethod - def is_eligible(cls, account: Account) -> bool: - if not dify_config.ENABLE_STEP_BY_STEP_TOUR: + def _is_eligible(self, account_started_at: datetime) -> bool: + if not self._enabled or self._rollout_started_at is None: return False - - rollout_started_at = dify_config.STEP_BY_STEP_TOUR_ROLLOUT_STARTED_AT - if rollout_started_at is None: - return False - - account_started_at = account.initialized_at or account.created_at - if account_started_at is None: - return False - - return ensure_naive_utc(account_started_at) >= ensure_naive_utc(rollout_started_at) - - @classmethod - def _get_state( - cls, - account_id: str, - *, - session: Session | scoped_session, - ) -> AccountStepByStepTourState | None: - stmt = select(AccountStepByStepTourState).where(AccountStepByStepTourState.account_id == account_id).limit(1) - return session.execute(stmt).scalar_one_or_none() - - @classmethod - def _ensure_state( - cls, - account_id: str, - *, - session: Session | scoped_session, - state: AccountStepByStepTourState | None, - ) -> AccountStepByStepTourState: - if state is None: - state = cls._get_state(account_id, session=session) - if state is not None: - return state - - state = AccountStepByStepTourState(account_id=account_id) - session.add(state) - try: - session.flush() - except IntegrityError: - # Another tab/device can create the account row between our read and insert. - session.rollback() - state = cls._get_state(account_id, session=session) - if state is None: - raise - return state + return ensure_naive_utc(account_started_at) >= ensure_naive_utc(self._rollout_started_at) @classmethod def _apply_action( cls, + state: StepByStepTourState, *, - state: AccountStepByStepTourState, - action: str, - task_id: str | None, - current_tenant_id: str, - ) -> None: - match action: + patch: StepByStepTourPatch, + workspace_id: str, + ) -> StepByStepTourState: + match patch.action: case "skip": - state.skipped = True - state.manually_enabled_workspace_ids = cls._remove_id( - state.manually_enabled_workspace_ids, - current_tenant_id, + return replace( + state, + skipped=True, + manually_enabled_workspace_ids=cls._remove_id( + state.manually_enabled_workspace_ids, + workspace_id, + ), ) case "complete_task": - if task_id is None: - raise ValueError("task_id is required") - cls._validate_task_id(task_id) - state.completed_task_ids = cls._add_id(state.completed_task_ids, task_id) + task_id = cls._require_task_id(patch.task_id) + return replace(state, completed_task_ids=cls._add_id(state.completed_task_ids, task_id)) case "uncomplete_task": - if task_id is None: - raise ValueError("task_id is required") - cls._validate_task_id(task_id) - state.completed_task_ids = cls._remove_id(state.completed_task_ids, task_id) + task_id = cls._require_task_id(patch.task_id) + return replace(state, completed_task_ids=cls._remove_id(state.completed_task_ids, task_id)) case "enable_current_workspace": - state.skipped = False - state.manually_enabled_workspace_ids = cls._add_id( - state.manually_enabled_workspace_ids, - current_tenant_id, - ) - state.manually_disabled_workspace_ids = cls._remove_id( - state.manually_disabled_workspace_ids, - current_tenant_id, + return replace( + state, + skipped=False, + manually_enabled_workspace_ids=cls._add_id( + state.manually_enabled_workspace_ids, + workspace_id, + ), + manually_disabled_workspace_ids=cls._remove_id( + state.manually_disabled_workspace_ids, + workspace_id, + ), ) case "disable_current_workspace": - state.manually_enabled_workspace_ids = cls._remove_id( - state.manually_enabled_workspace_ids, - current_tenant_id, - ) - state.manually_disabled_workspace_ids = cls._add_id( - state.manually_disabled_workspace_ids, - current_tenant_id, + return replace( + state, + manually_enabled_workspace_ids=cls._remove_id( + state.manually_enabled_workspace_ids, + workspace_id, + ), + manually_disabled_workspace_ids=cls._add_id( + state.manually_disabled_workspace_ids, + workspace_id, + ), ) case _: - raise ValueError(f"Unsupported action: {action}") - - @classmethod - def _build_response( - cls, - *, - state: AccountStepByStepTourState | None, - ) -> StepByStepTourStateResponse: - if state is None: - return { - "first_workspace_id": None, - "skipped": False, - "completed_task_ids": [], - "manually_enabled_workspace_ids": [], - "manually_disabled_workspace_ids": [], - "updated_at": None, - } - - return { - "first_workspace_id": state.first_workspace_id, - "skipped": state.skipped, - "completed_task_ids": cls._normalize_ids(state.completed_task_ids), - "manually_enabled_workspace_ids": cls._normalize_ids(state.manually_enabled_workspace_ids), - "manually_disabled_workspace_ids": cls._normalize_ids(state.manually_disabled_workspace_ids), - "updated_at": state.updated_at, - } + raise ValueError(f"Unsupported action: {patch.action}") @staticmethod - def _validate_task_id(task_id: str) -> None: - if task_id not in STEP_BY_STEP_TOUR_TASK_IDS: + def _require_workspace(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 + + @staticmethod + def _require_task_id(task_id: str | None) -> str: + if task_id is None: + raise ValueError("task_id is required") + if task_id not in _TASK_IDS: raise ValueError(f"Unsupported task_id: {task_id}") + return task_id @classmethod - def _add_id(cls, values: list[str], value: str) -> list[str]: + def _add_id(cls, values: tuple[str, ...], value: str) -> tuple[str, ...]: normalized = cls._normalize_ids(values) - if value in normalized: - return normalized - return [*normalized, value] + return normalized if value in normalized else (*normalized, value) @classmethod - def _remove_id(cls, values: list[str], value: str) -> list[str]: - return [item for item in cls._normalize_ids(values) if item != value] + def _remove_id(cls, values: tuple[str, ...], value: str) -> tuple[str, ...]: + return tuple(item for item in cls._normalize_ids(values) if item != value) @staticmethod - def _normalize_ids(values: list[str]) -> list[str]: - normalized: list[str] = [] - for value in values: - if value not in normalized: - normalized.append(value) - return normalized + def _normalize_ids(values: tuple[str, ...]) -> tuple[str, ...]: + return tuple(dict.fromkeys(values)) + + @staticmethod + def _to_result(state: StepByStepTourState | None) -> StepByStepTourResult: + if state is None: + return StepByStepTourResult() + return StepByStepTourResult( + first_workspace_id=state.first_workspace_id, + skipped=state.skipped, + completed_task_ids=tuple(dict.fromkeys(state.completed_task_ids)), + manually_enabled_workspace_ids=tuple(dict.fromkeys(state.manually_enabled_workspace_ids)), + manually_disabled_workspace_ids=tuple(dict.fromkeys(state.manually_disabled_workspace_ids)), + updated_at=state.updated_at, + ) diff --git a/api/tests/unit_tests/commands/test_generate_swagger_markdown_docs.py b/api/tests/unit_tests/commands/test_generate_swagger_markdown_docs.py index cff6695e414..9231e274d5c 100644 --- a/api/tests/unit_tests/commands/test_generate_swagger_markdown_docs.py +++ b/api/tests/unit_tests/commands/test_generate_swagger_markdown_docs.py @@ -238,6 +238,42 @@ def test_patch_union_schema_markdown_fills_regular_schema_union_property(tmp_pat assert "| value | string
integer
number
boolean | | No |" in patched +def test_patch_union_schema_markdown_preserves_nullable_enum_values(tmp_path: Path): + module = _load_generate_swagger_markdown_docs_module() + spec_path = tmp_path / "console-openapi.json" + spec_path.write_text( + json.dumps( + { + "components": { + "schemas": { + "StepByStepTourStatePatchPayload": { + "properties": { + "task_id": { + "anyOf": [ + {"enum": ["home", "studio"], "type": "string"}, + {"type": "null"}, + ], + }, + }, + }, + }, + } + } + ), + encoding="utf-8", + ) + markdown = """#### StepByStepTourStatePatchPayload + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| task_id | string | Task ID | No | +""" + + patched = module._patch_union_schema_markdown(markdown, spec_path) + + assert '| task_id | string,
**Available values:** "home", "studio" | Task ID | No |' in patched + + def test_patch_union_schema_markdown_fills_array_item_union_property(tmp_path: Path): module = _load_generate_swagger_markdown_docs_module() spec_path = tmp_path / "console-openapi.json" diff --git a/api/tests/unit_tests/controllers/console/test_notification.py b/api/tests/unit_tests/controllers/console/test_notification.py new file mode 100644 index 00000000000..48843d1af8a --- /dev/null +++ b/api/tests/unit_tests/controllers/console/test_notification.py @@ -0,0 +1,77 @@ +from inspect import unwrap +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from controllers.console.notification import ( + DismissNotificationPayload, + NotificationApi, + NotificationDismissApi, +) +from machinery.context import RequestContext +from services.entities.notification_entities import NotificationItem, NotificationResult + + +def _request_context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id="workspace-1", + ) + + +def test_get_notification_delegates_and_serializes_result() -> None: + service = Mock() + service.get_active.return_value = NotificationResult( + should_show=True, + notifications=( + NotificationItem( + notification_id="notification-1", + frequency="once", + lang="en-US", + title="Title", + subtitle="Subtitle", + body="Body", + title_pic_url="https://example.com/title.png", + ), + ), + ) + services = SimpleNamespace(notifications=service) + api = NotificationApi() + method = unwrap(api.get) + context = _request_context() + + with patch("controllers.console.notification.application_services", return_value=services): + result, status = method(api, context) + + assert status == 200 + assert result == { + "should_show": True, + "notifications": [ + { + "notification_id": "notification-1", + "frequency": "once", + "lang": "en-US", + "title": "Title", + "subtitle": "Subtitle", + "body": "Body", + "title_pic_url": "https://example.com/title.png", + } + ], + } + service.get_active.assert_called_once_with(context) + + +def test_dismiss_notification_delegates_with_stable_account_context() -> None: + service = Mock() + services = SimpleNamespace(notifications=service) + api = NotificationDismissApi() + method = unwrap(api.post) + context = _request_context() + + with patch("controllers.console.notification.application_services", return_value=services): + result, status = method(api, DismissNotificationPayload(notification_id="notification-1"), context) + + assert status == 200 + assert result == {"result": "success"} + service.dismiss.assert_called_once_with(context, "notification-1") diff --git a/api/tests/unit_tests/controllers/console/test_onboarding.py b/api/tests/unit_tests/controllers/console/test_onboarding.py index 8d613f7c202..90a9521a2fa 100644 --- a/api/tests/unit_tests/controllers/console/test_onboarding.py +++ b/api/tests/unit_tests/controllers/console/test_onboarding.py @@ -2,47 +2,48 @@ from __future__ import annotations from datetime import UTC, datetime from inspect import unwrap -from unittest.mock import Mock +from types import SimpleNamespace +from unittest.mock import Mock, patch import pytest -from flask import Flask from pydantic import ValidationError from controllers.console.onboarding import ( StepByStepTourStateApi, StepByStepTourStatePatchPayload, + StepByStepTourStateResponse, ) -from extensions.ext_database import db -from models.account import Account, AccountStatus -from services.step_by_step_tour_service import StepByStepTourService +from machinery.context import RequestContext +from services.entities.onboarding_entities import StepByStepTourPatch, StepByStepTourResult -def _account() -> Account: - account = Account(name="User", email="user@example.com", status=AccountStatus.ACTIVE) - account.id = "account-1" - return account +def _request_context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id="workspace-1", + ) -def _state_response() -> dict[str, object]: - return { - "first_workspace_id": "workspace-1", - "skipped": False, - "completed_task_ids": ["home"], - "manually_enabled_workspace_ids": [], - "manually_disabled_workspace_ids": [], - "updated_at": datetime(2026, 6, 28, tzinfo=UTC), - } +def _state_result() -> StepByStepTourResult: + return StepByStepTourResult( + first_workspace_id="workspace-1", + completed_task_ids=("home",), + updated_at=datetime(2026, 6, 28, tzinfo=UTC), + ) -def test_get_step_by_step_tour_state(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - get_state = Mock(return_value=_state_response()) - monkeypatch.setattr(StepByStepTourService, "get_state", get_state) - +def test_get_step_by_step_tour_state_delegates_with_request_context() -> None: + service = Mock() + service.get_state.return_value = _state_result() + services = SimpleNamespace(step_by_step_tour=service) api = StepByStepTourStateApi() method = unwrap(api.get) + context = _request_context() - with app.test_request_context("/console/api/onboarding/step-by-step-tour/state", method="GET"): - result = method(api, "workspace-1", _account()) + with patch("controllers.console.onboarding.application_services", return_value=services): + result = method(api, context) assert result == { "first_workspace_id": "workspace-1", @@ -52,35 +53,26 @@ def test_get_step_by_step_tour_state(app: Flask, monkeypatch: pytest.MonkeyPatch "manually_disabled_workspace_ids": [], "updated_at": "2026-06-28T00:00:00Z", } - get_state.assert_called_once() - assert get_state.call_args.kwargs["current_tenant_id"] == "workspace-1" - assert get_state.call_args.kwargs["session"] is db.session + service.get_state.assert_called_once_with(context) -def test_patch_step_by_step_tour_state_passes_action_payload( - app: Flask, - monkeypatch: pytest.MonkeyPatch, -) -> None: - patch_state = Mock(return_value=_state_response()) - monkeypatch.setattr(StepByStepTourService, "patch_state", patch_state) - +def test_patch_step_by_step_tour_state_maps_transport_payload_to_command() -> None: + service = Mock() + service.patch_state.return_value = _state_result() + services = SimpleNamespace(step_by_step_tour=service) api = StepByStepTourStateApi() method = unwrap(api.patch) - payload = {"action": "complete_task", "task_id": "studio"} + context = _request_context() + payload = StepByStepTourStatePatchPayload.model_validate({"action": "complete_task", "task_id": "studio"}) - req_data = StepByStepTourStatePatchPayload.model_validate(payload) - with app.test_request_context( - "/console/api/onboarding/step-by-step-tour/state", - method="PATCH", - json=payload, - ): - result = method(api, req_data, "workspace-1", _account()) + with patch("controllers.console.onboarding.application_services", return_value=services): + result = method(api, payload, context) assert result["completed_task_ids"] == ["home"] - patch_state.assert_called_once() - assert patch_state.call_args.kwargs["current_tenant_id"] == "workspace-1" - assert patch_state.call_args.kwargs["patch"] == payload - assert patch_state.call_args.kwargs["session"] is db.session + service.patch_state.assert_called_once_with( + context, + StepByStepTourPatch(action="complete_task", task_id="studio"), + ) def test_patch_payload_rejects_non_action_fields() -> None: @@ -96,3 +88,21 @@ def test_patch_payload_rejects_task_id_without_task_action() -> None: def test_patch_payload_requires_action() -> None: with pytest.raises(ValidationError): StepByStepTourStatePatchPayload.model_validate({"task_id": "home"}) + + +def test_step_by_step_tour_schemas_preserve_enum_values() -> None: + patch_schema = StepByStepTourStatePatchPayload.model_json_schema() + action_schema = patch_schema["properties"]["action"] + task_id_schema = patch_schema["properties"]["task_id"] + task_id_values = next(candidate["enum"] for candidate in task_id_schema["anyOf"] if "enum" in candidate) + response_schema = StepByStepTourStateResponse.model_json_schema() + + assert set(action_schema["enum"]) == { + "skip", + "complete_task", + "uncomplete_task", + "enable_current_workspace", + "disable_current_workspace", + } + assert set(task_id_values) == {"home", "studio", "knowledge", "integration"} + assert set(response_schema["properties"]["completed_task_ids"]["items"]["enum"]) == set(task_id_values) 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 ee3d2125204..191f13cf7e3 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -382,6 +382,8 @@ def test_build_application_services_wires_account_profile_repository( assert email_registration._registration._session_factory is sqlite_session_factory assert services.accounts.education._accounts is accounts assert services.accounts.deletion._accounts is accounts + assert services.notifications._accounts is accounts + assert services.step_by_step_tour._accounts is accounts assert services.accounts.deletion._memberships is services.workspace_queries._workspaces integrations = services.accounts.integrations._integrations assert isinstance(integrations, SQLAlchemyAccountIntegrationRepository) diff --git a/api/tests/unit_tests/repositories/test_step_by_step_tour_repository.py b/api/tests/unit_tests/repositories/test_step_by_step_tour_repository.py new file mode 100644 index 00000000000..a6439f58bc5 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_step_by_step_tour_repository.py @@ -0,0 +1,171 @@ +from contextlib import nullcontext +from dataclasses import replace +from datetime import datetime +from typing import cast +from unittest.mock import MagicMock, Mock + +import pytest +from sqlalchemy.exc import IntegrityError, OperationalError +from sqlalchemy.orm import Session, sessionmaker + +from models.onboarding import AccountStepByStepTourState +from repositories.step_by_step_tour_repository import ( + SQLAlchemyStepByStepTourStateRepository, + _is_retryable_mysql_lock_error, +) + + +class _ErrnoOnlyError(Exception): + def __init__(self, errno: int | str) -> None: + super().__init__() + self.errno = errno + + +def test_mutate_creates_and_updates_state_in_repository_owned_transaction( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyStepByStepTourStateRepository(sqlite_session_factory) + + saved = repository.mutate( + "account-1", + lambda state: replace(state, completed_task_ids=("home",)), + ) + reloaded = repository.get("account-1") + + assert saved.first_workspace_id is None + assert saved.completed_task_ids == ("home",) + assert saved.updated_at is not None + assert reloaded == saved + + +def test_initialize_creates_state_with_first_workspace_atomically( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyStepByStepTourStateRepository(sqlite_session_factory) + + result = repository.initialize("account-1", "workspace-1") + + assert result.first_workspace_id == "workspace-1" + assert repository.get("account-1") == result + + +def test_initialize_claims_empty_state_once_without_overwriting_winner( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyStepByStepTourStateRepository(sqlite_session_factory) + with sqlite_session_factory() as session: + session.add(AccountStepByStepTourState(account_id="account-1")) + session.commit() + + first = repository.initialize("account-1", "workspace-1") + second = repository.initialize("account-1", "workspace-2") + + assert first.first_workspace_id == "workspace-1" + assert second.first_workspace_id == "workspace-1" + + +def test_mutate_cannot_clear_or_overwrite_first_workspace( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyStepByStepTourStateRepository(sqlite_session_factory) + repository.initialize("account-1", "workspace-1") + + result = repository.mutate( + "account-1", + lambda state: replace(state, first_workspace_id="workspace-2", skipped=True), + ) + + assert result.first_workspace_id == "workspace-1" + assert result.skipped is True + + +def test_sequential_mutations_replay_against_latest_state( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyStepByStepTourStateRepository(sqlite_session_factory) + + repository.mutate("account-1", lambda state: replace(state, completed_task_ids=("home",))) + result = repository.mutate( + "account-1", + lambda state: replace(state, completed_task_ids=(*state.completed_task_ids, "studio")), + ) + + assert result.completed_task_ids == ("home", "studio") + + +def test_mutate_replays_after_concurrent_create_conflict() -> None: + concurrent_state = AccountStepByStepTourState(account_id="account-1") + concurrent_state.completed_task_ids = ["home"] + concurrent_state.updated_at = datetime(2026, 8, 13) + session = MagicMock(spec=Session) + session.execute.return_value.scalar_one_or_none.side_effect = [None, concurrent_state] + session.flush.side_effect = IntegrityError("insert", {}, Exception("duplicate")) + factory = cast(sessionmaker[Session], Mock(return_value=nullcontext(session))) + repository = SQLAlchemyStepByStepTourStateRepository(factory) + + result = repository.mutate( + "account-1", + lambda state: replace(state, completed_task_ids=(*state.completed_task_ids, "studio")), + ) + + assert result.completed_task_ids == ("home", "studio") + session.rollback.assert_called_once_with() + initial_probe = session.execute.call_args_list[0].args[0] + replay_statement = session.execute.call_args_list[1].args[0] + assert initial_probe._for_update_arg is None + assert replay_statement._for_update_arg is not None + + +def test_mutate_retries_mysql_deadlock_with_fresh_session() -> None: + concurrent_state = AccountStepByStepTourState(account_id="account-1") + concurrent_state.completed_task_ids = ["home"] + concurrent_state.updated_at = datetime(2026, 8, 13) + + deadlocked_session = MagicMock(spec=Session) + deadlocked_session.execute.return_value.scalar_one_or_none.return_value = None + deadlocked_session.flush.side_effect = OperationalError( + "INSERT", + {}, + Exception(1213, "Deadlock found when trying to get lock"), + ) + + retry_session = MagicMock(spec=Session) + retry_session.execute.return_value.scalar_one_or_none.side_effect = [concurrent_state, concurrent_state] + factory = Mock(side_effect=[nullcontext(deadlocked_session), nullcontext(retry_session)]) + repository = SQLAlchemyStepByStepTourStateRepository(cast(sessionmaker[Session], factory)) + + result = repository.mutate( + "account-1", + lambda state: replace(state, completed_task_ids=(*state.completed_task_ids, "studio")), + ) + + assert result.completed_task_ids == ("home", "studio") + assert factory.call_count == 2 + retry_lock_statement = retry_session.execute.call_args_list[1].args[0] + assert retry_lock_statement._for_update_arg is not None + + +@pytest.mark.parametrize( + ("orig", "expected"), + [ + pytest.param(_ErrnoOnlyError(1205), True, id="errno-attribute"), + pytest.param(Exception(1213, "deadlock"), True, id="integer-args-code"), + pytest.param(Exception("1213", "deadlock"), True, id="string-args-code"), + pytest.param(Exception(9999, "other error"), False, id="non-retryable-code"), + pytest.param(Exception(True), False, id="boolean-is-not-an-error-code"), + pytest.param(Exception(), False, id="missing-error-code"), + ], +) +def test_mysql_lock_error_detection_preserves_errno_and_args_coverage( + orig: BaseException, + expected: bool, +) -> None: + exc = OperationalError("statement", {}, orig) + + assert _is_retryable_mysql_lock_error(exc) is expected + + +def test_get_returns_none_for_unknown_account( + sqlite_session_factory: sessionmaker[Session], +) -> None: + assert SQLAlchemyStepByStepTourStateRepository(sqlite_session_factory).get("missing") is None diff --git a/api/tests/unit_tests/services/test_notification_gateway.py b/api/tests/unit_tests/services/test_notification_gateway.py new file mode 100644 index 00000000000..9df67e7a325 --- /dev/null +++ b/api/tests/unit_tests/services/test_notification_gateway.py @@ -0,0 +1,63 @@ +from unittest.mock import patch + +from services.entities.notification_entities import NotificationContent +from services.notification_gateway import BillingNotificationGateway + + +def test_get_active_maps_billing_proto_json_contract() -> None: + payload = { + "shouldShow": True, + "notifications": [ + { + "notificationId": "notification-1", + "frequency": "once", + "contents": { + "en-US": { + "lang": "en-US", + "title": "Title", + "subtitle": "Subtitle", + "body": "Body", + "titlePicUrl": "title.png", + } + }, + } + ], + } + + with patch("services.notification_gateway.BillingService.get_account_notification", return_value=payload): + result = BillingNotificationGateway().get_active("account-1") + + assert result.should_show is True + assert result.notifications[0].notification_id == "notification-1" + assert result.notifications[0].contents["en-US"].title_pic_url == "title.png" + + +def test_get_active_omits_empty_localized_content_so_service_can_fall_back() -> None: + empty_localized_content: dict[str, str] = {} + payload = { + "shouldShow": True, + "notifications": [ + { + "notificationId": "notification-1", + "frequency": "once", + "contents": { + "zh-Hans": empty_localized_content, + "en-US": {"lang": "en-US", "title": "Title"}, + }, + } + ], + } + + with patch("services.notification_gateway.BillingService.get_account_notification", return_value=payload): + result = BillingNotificationGateway().get_active("account-1") + + assert result.notifications[0].contents == { + "en-US": NotificationContent("en-US", "Title", "", "", ""), + } + + +def test_dismiss_delegates_to_billing_service() -> None: + with patch("services.notification_gateway.BillingService.dismiss_notification") as dismiss: + BillingNotificationGateway().dismiss("notification-1", "account-1") + + dismiss.assert_called_once_with(notification_id="notification-1", account_id="account-1") diff --git a/api/tests/unit_tests/services/test_notification_service.py b/api/tests/unit_tests/services/test_notification_service.py new file mode 100644 index 00000000000..3be7f08a6f7 --- /dev/null +++ b/api/tests/unit_tests/services/test_notification_service.py @@ -0,0 +1,138 @@ +from datetime import datetime +from unittest.mock import Mock + +import pytest + +from machinery.context import RequestContext +from services.account_ports import AccountRepository +from services.entities.account_entities import AccountSnapshot +from services.entities.notification_entities import ( + AccountNotification, + AccountNotificationBatch, + NotificationContent, + NotificationItem, + NotificationResult, +) +from services.notification_service import NotificationService + + +def _context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id="workspace-1", + ) + + +class NotificationGatewayStub: + def __init__(self, batch: AccountNotificationBatch) -> None: + self.batch = batch + self.get_account_ids: list[str] = [] + self.dismissals: list[tuple[str, str]] = [] + + def get_active(self, account_id: str) -> AccountNotificationBatch: + self.get_account_ids.append(account_id) + return self.batch + + def dismiss(self, notification_id: str, account_id: str) -> None: + self.dismissals.append((notification_id, account_id)) + + +def _account(language: str | None = "zh-Hans") -> AccountSnapshot: + return AccountSnapshot( + id="account-1", + name="Account", + email="account@example.com", + avatar=None, + is_password_set=False, + interface_language=language, + 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 _accounts(account: AccountSnapshot | None) -> Mock: + accounts = Mock(spec=AccountRepository) + accounts.get.return_value = account + return accounts + + +def _notification(contents: dict[str, NotificationContent]) -> AccountNotification: + return AccountNotification( + notification_id="notification-1", + frequency="once", + contents=contents, + ) + + +def test_get_active_localizes_notification_for_account_language() -> None: + chinese = NotificationContent("zh-Hans", "标题", "副标题", "正文", "zh.png") + english = NotificationContent("en-US", "Title", "Subtitle", "Body", "en.png") + gateway = NotificationGatewayStub( + AccountNotificationBatch(True, (_notification({"zh-Hans": chinese, "en-US": english}),)) + ) + service = NotificationService(accounts=_accounts(_account()), notifications=gateway) + + result = service.get_active(_context()) + + assert result == NotificationResult( + should_show=True, + notifications=(NotificationItem("notification-1", "once", "zh-Hans", "标题", "副标题", "正文", "zh.png"),), + ) + assert gateway.get_account_ids == ["account-1"] + + +def test_get_active_falls_back_to_english() -> None: + english = NotificationContent("en-US", "Title", "Subtitle", "Body", "en.png") + gateway = NotificationGatewayStub(AccountNotificationBatch(True, (_notification({"en-US": english}),))) + service = NotificationService(accounts=_accounts(_account("fr-FR")), notifications=gateway) + + result = service.get_active(_context()) + + assert result.notifications[0].lang == "en-US" + assert result.notifications[0].title == "Title" + + +def test_get_active_skips_account_query_when_gateway_says_not_to_show() -> None: + accounts = _accounts(None) + service = NotificationService( + accounts=accounts, + notifications=NotificationGatewayStub(AccountNotificationBatch(False, ())), + ) + + result = service.get_active(_context()) + + assert result == NotificationResult(False, ()) + accounts.get.assert_not_called() + + +def test_get_active_uses_empty_content_when_notification_has_no_translations() -> None: + gateway = NotificationGatewayStub(AccountNotificationBatch(True, (_notification({}),))) + service = NotificationService(accounts=_accounts(_account(None)), notifications=gateway) + + result = service.get_active(_context()) + + assert result.notifications == (NotificationItem("notification-1", "once", "en-US", "", "", "", ""),) + + +def test_get_active_rejects_unknown_admitted_account() -> None: + gateway = NotificationGatewayStub(AccountNotificationBatch(True, (_notification({}),))) + service = NotificationService(accounts=_accounts(None), notifications=gateway) + + with pytest.raises(RuntimeError, match="unknown account"): + service.get_active(_context()) + + +def test_dismiss_delegates_identifiers_to_gateway() -> None: + gateway = NotificationGatewayStub(AccountNotificationBatch(False, ())) + service = NotificationService(accounts=_accounts(_account()), notifications=gateway) + + service.dismiss(_context(), "notification-1") + + assert gateway.dismissals == [("notification-1", "account-1")] diff --git a/api/tests/unit_tests/services/test_step_by_step_tour_service.py b/api/tests/unit_tests/services/test_step_by_step_tour_service.py index 7a99fa61444..40017bb7798 100644 --- a/api/tests/unit_tests/services/test_step_by_step_tour_service.py +++ b/api/tests/unit_tests/services/test_step_by_step_tour_service.py @@ -1,230 +1,213 @@ from __future__ import annotations -from datetime import UTC, datetime +from collections.abc import Callable +from dataclasses import replace +from datetime import datetime +from unittest.mock import Mock import pytest -from sqlalchemy import event, select -from sqlalchemy.orm import Session, sessionmaker -from enums import DeploymentEdition -from models.account import Account, AccountStatus -from models.onboarding import AccountStepByStepTourState +from machinery.context import RequestContext +from services.account_ports import AccountRepository +from services.entities.account_entities import AccountSnapshot +from services.entities.onboarding_entities import StepByStepTourPatch, StepByStepTourResult, StepByStepTourState from services.step_by_step_tour_service import StepByStepTourService -from tests.unit_tests.config_override import apply_config_overrides -def _account(*, initialized_at: datetime | None = None, created_at: datetime | None = None) -> Account: - account = Account(name="User", email="user@example.com", status=AccountStatus.ACTIVE) - account.id = "account-1" - account.initialized_at = initialized_at - account.created_at = created_at or datetime(2026, 6, 28) - return account - - -def _state() -> AccountStepByStepTourState: - state = AccountStepByStepTourState(account_id="account-1") - state.updated_at = datetime(2026, 6, 28, tzinfo=UTC) - return state - - -def _persist_state(session: Session, state: AccountStepByStepTourState) -> None: - session.add(state) - session.commit() - - -def _load_state(session: Session) -> AccountStepByStepTourState | None: - return session.scalar( - select(AccountStepByStepTourState).where(AccountStepByStepTourState.account_id == "account-1") +def _context(*, workspace_id: str | None = "workspace-1") -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id=workspace_id, ) -def _set_tour_config(monkeypatch: pytest.MonkeyPatch, *, enabled: bool, rollout_started_at: datetime | None) -> None: - apply_config_overrides( - monkeypatch, - ENABLE_STEP_BY_STEP_TOUR=enabled, - STEP_BY_STEP_TOUR_ROLLOUT_STARTED_AT=rollout_started_at, +class StateRepositoryStub: + def __init__(self, state: StepByStepTourState | None = None) -> None: + self.state = state + self.get_account_ids: list[str] = [] + self.initialize_calls: list[tuple[str, str]] = [] + self.mutation_account_ids: list[str] = [] + + def get(self, account_id: str) -> StepByStepTourState | None: + self.get_account_ids.append(account_id) + return self.state + + def initialize(self, account_id: str, first_workspace_id: str) -> StepByStepTourState: + self.initialize_calls.append((account_id, first_workspace_id)) + if self.state is None: + self.state = StepByStepTourState(account_id=account_id, first_workspace_id=first_workspace_id) + elif self.state.first_workspace_id is None: + self.state = replace(self.state, first_workspace_id=first_workspace_id) + return self.state + + def mutate( + self, + account_id: str, + mutation: Callable[[StepByStepTourState], StepByStepTourState], + ) -> StepByStepTourState: + self.mutation_account_ids.append(account_id) + if self.state is None: + self.state = StepByStepTourState(account_id=account_id) + self.state = mutation(self.state) + return self.state + + +def _account(*, started_at: datetime = datetime(2026, 6, 28)) -> 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=started_at, + created_at=started_at, ) -def test_get_state_creates_state_and_records_first_workspace_for_eligible_account( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - _set_tour_config(monkeypatch, enabled=True, rollout_started_at=datetime(2026, 6, 1)) +def _accounts(account: AccountSnapshot | None) -> Mock: + accounts = Mock(spec=AccountRepository) + accounts.get.return_value = account + return accounts - result = StepByStepTourService.get_state( - account=_account(initialized_at=datetime(2026, 6, 28)), - current_tenant_id="workspace-1", - session=sqlite_session, + +def _service( + *, + states: StateRepositoryStub, + account: AccountSnapshot | None = None, + enabled: bool = True, + rollout_started_at: datetime | None = datetime(2026, 6, 1), +) -> StepByStepTourService: + return StepByStepTourService( + accounts=_accounts(account or _account()), + states=states, + enabled=enabled, + rollout_started_at=rollout_started_at, ) - assert result["first_workspace_id"] == "workspace-1" - assert result["completed_task_ids"] == [] - with sqlite_session_factory() as observer: - persisted = _load_state(observer) - assert persisted is not None - assert persisted.account_id == "account-1" - assert persisted.first_workspace_id == "workspace-1" + +def test_get_state_creates_state_and_records_first_workspace_for_eligible_account() -> None: + states = StateRepositoryStub() + + result = _service(states=states).get_state(_context()) + + assert result.first_workspace_id == "workspace-1" + assert states.get_account_ids == [] + assert states.initialize_calls == [("account-1", "workspace-1")] + assert states.mutation_account_ids == [] -def test_is_eligible_does_not_depend_on_cloud_edition(monkeypatch: pytest.MonkeyPatch) -> None: - _set_tour_config(monkeypatch, enabled=True, rollout_started_at=datetime(2026, 6, 1)) - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) +def test_get_state_returns_existing_state_without_rewriting_first_workspace() -> None: + state = StepByStepTourState(account_id="account-1", first_workspace_id="workspace-original") + states = StateRepositoryStub(state) - result = StepByStepTourService.is_eligible(_account(initialized_at=datetime(2026, 6, 28))) + result = _service(states=states).get_state(_context(workspace_id="workspace-current")) - assert result is True + assert result.first_workspace_id == "workspace-original" + assert states.initialize_calls == [("account-1", "workspace-current")] + assert states.mutation_account_ids == [] -def test_get_state_does_not_create_state_for_ineligible_account_without_existing_state( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - _set_tour_config(monkeypatch, enabled=True, rollout_started_at=datetime(2026, 6, 1)) +def test_get_state_does_not_create_state_for_ineligible_account() -> None: + states = StateRepositoryStub() + service = _service(states=states, account=_account(started_at=datetime(2026, 5, 31))) - result = StepByStepTourService.get_state( - account=_account(initialized_at=datetime(2026, 5, 31)), - current_tenant_id="workspace-1", - session=sqlite_session, + result = service.get_state(_context()) + + assert result == StepByStepTourResult() + assert states.get_account_ids == ["account-1"] + assert states.mutation_account_ids == [] + + +def test_get_state_does_not_create_state_when_tour_is_disabled() -> None: + states = StateRepositoryStub() + + result = _service(states=states, enabled=False).get_state(_context()) + + assert result == StepByStepTourResult() + assert states.get_account_ids == ["account-1"] + + +def test_patch_state_persists_even_when_tour_is_disabled() -> None: + states = StateRepositoryStub() + service = _service(states=states, enabled=False) + + result = service.patch_state(_context(workspace_id="workspace-2"), StepByStepTourPatch("enable_current_workspace")) + + assert result.manually_enabled_workspace_ids == ("workspace-2",) + assert states.mutation_account_ids == ["account-1"] + + +def test_patch_state_skip_removes_current_workspace_enable() -> None: + states = StateRepositoryStub( + StepByStepTourState( + account_id="account-1", + manually_enabled_workspace_ids=("workspace-1", "workspace-2"), + ) ) - assert result == { - "first_workspace_id": None, - "skipped": False, - "completed_task_ids": [], - "manually_enabled_workspace_ids": [], - "manually_disabled_workspace_ids": [], - "updated_at": None, - } - with sqlite_session_factory() as observer: - assert _load_state(observer) is None + result = _service(states=states).patch_state(_context(), StepByStepTourPatch("skip")) + + assert result.skipped is True + assert result.manually_enabled_workspace_ids == ("workspace-2",) -def test_patch_state_persists_even_when_account_is_not_eligible( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - _set_tour_config(monkeypatch, enabled=False, rollout_started_at=datetime(2026, 6, 1)) - - result = StepByStepTourService.patch_state( - account=_account(initialized_at=datetime(2026, 6, 28)), - current_tenant_id="workspace-2", - patch={"action": "enable_current_workspace"}, - session=sqlite_session, +def test_patch_state_disable_moves_current_workspace_to_disabled() -> None: + states = StateRepositoryStub( + StepByStepTourState( + account_id="account-1", + manually_enabled_workspace_ids=("workspace-1", "workspace-2"), + ) ) - assert result["skipped"] is False - assert result["manually_enabled_workspace_ids"] == ["workspace-2"] - assert result["manually_disabled_workspace_ids"] == [] - with sqlite_session_factory() as observer: - persisted = _load_state(observer) - assert persisted is not None - assert persisted.manually_enabled_workspace_ids == ["workspace-2"] - - -def test_patch_state_skip_action_sets_skipped_and_removes_current_workspace_enable( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - _set_tour_config(monkeypatch, enabled=False, rollout_started_at=datetime(2026, 6, 1)) - state = _state() - state.manually_enabled_workspace_ids = ["workspace-1", "workspace-2"] - _persist_state(sqlite_session, state) - - result = StepByStepTourService.patch_state( - account=_account(initialized_at=datetime(2026, 6, 28)), - current_tenant_id="workspace-1", - patch={"action": "skip"}, - session=sqlite_session, + result = _service(states=states).patch_state( + _context(), + StepByStepTourPatch("disable_current_workspace"), ) - assert result["skipped"] is True - assert result["manually_enabled_workspace_ids"] == ["workspace-2"] - assert result["manually_disabled_workspace_ids"] == [] - assert _load_state(sqlite_session) is state + assert result.manually_enabled_workspace_ids == ("workspace-2",) + assert result.manually_disabled_workspace_ids == ("workspace-1",) -def test_patch_state_disable_action_moves_current_workspace_to_disabled( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - _set_tour_config(monkeypatch, enabled=False, rollout_started_at=datetime(2026, 6, 1)) - state = _state() - state.manually_enabled_workspace_ids = ["workspace-1", "workspace-2"] - _persist_state(sqlite_session, state) +def test_patch_state_complete_and_uncomplete_task() -> None: + states = StateRepositoryStub(StepByStepTourState(account_id="account-1", completed_task_ids=("home",))) + service = _service(states=states) - result = StepByStepTourService.patch_state( - account=_account(initialized_at=datetime(2026, 6, 28)), - current_tenant_id="workspace-1", - patch={"action": "disable_current_workspace"}, - session=sqlite_session, + service.patch_state(_context(), StepByStepTourPatch("complete_task", "studio")) + result = service.patch_state(_context(), StepByStepTourPatch("uncomplete_task", "home")) + + assert result.completed_task_ids == ("studio",) + + +def test_rejects_unsupported_task_id() -> None: + with pytest.raises(ValueError, match="Unsupported task_id"): + StepByStepTourService._require_task_id("unknown") + + +def test_rejects_missing_workspace_before_using_state_repository() -> None: + states = StateRepositoryStub() + + with pytest.raises(RuntimeError, match="did not resolve an active workspace"): + _service(states=states).patch_state(_context(workspace_id=None), StepByStepTourPatch("skip")) + + assert states.mutation_account_ids == [] + + +def test_get_state_rejects_unknown_admitted_account() -> None: + states = StateRepositoryStub() + service = StepByStepTourService( + accounts=_accounts(None), + states=states, + enabled=True, + rollout_started_at=datetime(2026, 6, 1), ) - assert result["manually_enabled_workspace_ids"] == ["workspace-2"] - assert result["manually_disabled_workspace_ids"] == ["workspace-1"] - assert _load_state(sqlite_session) is state - - -def test_patch_state_complete_and_uncomplete_task( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - _set_tour_config(monkeypatch, enabled=False, rollout_started_at=datetime(2026, 6, 1)) - state = _state() - state.completed_task_ids = ["home"] - _persist_state(sqlite_session, state) - - StepByStepTourService.patch_state( - account=_account(initialized_at=datetime(2026, 6, 28)), - current_tenant_id="workspace-1", - patch={"action": "complete_task", "task_id": "studio"}, - session=sqlite_session, - ) - result = StepByStepTourService.patch_state( - account=_account(initialized_at=datetime(2026, 6, 28)), - current_tenant_id="workspace-1", - patch={"action": "uncomplete_task", "task_id": "home"}, - session=sqlite_session, - ) - - assert result["completed_task_ids"] == ["studio"] - - -def test_patch_state_recovers_when_concurrent_request_created_state( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - _set_tour_config(monkeypatch, enabled=False, rollout_started_at=datetime(2026, 6, 1)) - existing_state = _state() - existing_state.manually_enabled_workspace_ids = ["workspace-1"] - lifecycle_events: list[str] = [] - - @event.listens_for(sqlite_session, "before_flush", once=True) - def add_conflicting_pending_state(session: Session, _flush_context, _instances) -> None: - lifecycle_events.append("before_flush") - session.add(AccountStepByStepTourState(account_id="account-1")) - - @event.listens_for(sqlite_session, "after_soft_rollback", once=True) - def persist_winning_request(_session: Session, _previous_transaction) -> None: - lifecycle_events.append("after_soft_rollback") - with sqlite_session_factory() as winner: - winner.add(existing_state) - winner.commit() - - result = StepByStepTourService.patch_state( - account=_account(initialized_at=datetime(2026, 6, 28)), - current_tenant_id="workspace-2", - patch={"action": "enable_current_workspace"}, - session=sqlite_session, - ) - - assert result["manually_enabled_workspace_ids"] == ["workspace-1", "workspace-2"] - assert lifecycle_events == ["before_flush", "after_soft_rollback"] - with sqlite_session_factory() as observer: - persisted = _load_state(observer) - assert persisted is not None - assert persisted.manually_enabled_workspace_ids == ["workspace-1", "workspace-2"] + with pytest.raises(RuntimeError, match="unknown account"): + service.get_state(_context()) From 5a5b9c6fec1cd2c26def10e3a6ae0d384da361ba Mon Sep 17 00:00:00 2001 From: yyh <92089059+lyzno1@users.noreply.github.com> Date: Mon, 31 Aug 2026 05:46:55 +0000 Subject: [PATCH 09/21] fix(web): make plugin detail drawer non-modal (#41473) --- web/app/components/plugins/plugin-detail-panel/index.tsx | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/web/app/components/plugins/plugin-detail-panel/index.tsx b/web/app/components/plugins/plugin-detail-panel/index.tsx index 22f5a212cab..094898ea58e 100644 --- a/web/app/components/plugins/plugin-detail-panel/index.tsx +++ b/web/app/components/plugins/plugin-detail-panel/index.tsx @@ -4,7 +4,6 @@ import type { PluginDetail } from '@/app/components/plugins/types' import { cn } from '@langgenius/dify-ui/cn' import { Drawer, - DrawerBackdrop, DrawerContent, DrawerPopup, DrawerPortal, @@ -68,18 +67,18 @@ const PluginDetailPanel: FC = ({ return ( { if (!open) onHide() }} > - - + From 02222e8fcebca81260dbe20b72b2bc3d3c571f9c Mon Sep 17 00:00:00 2001 From: Jingyi Date: Mon, 31 Aug 2026 05:50:29 +0000 Subject: [PATCH 10/21] fix(web): preserve plugin footer text layout (#41514) --- web/app/components/integrations/tool-provider-card.tsx | 4 ++-- web/app/components/plugins/card/base/org-info.tsx | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/web/app/components/integrations/tool-provider-card.tsx b/web/app/components/integrations/tool-provider-card.tsx index f6aca486a20..27dc904c968 100644 --- a/web/app/components/integrations/tool-provider-card.tsx +++ b/web/app/components/integrations/tool-provider-card.tsx @@ -131,13 +131,13 @@ function IntegrationsToolProviderCard({
{!!org && ( <> -
+
{org}
/
)} -
+
{name}
diff --git a/web/app/components/plugins/card/base/org-info.tsx b/web/app/components/plugins/card/base/org-info.tsx index 2d9ef6ce036..9989cd454cb 100644 --- a/web/app/components/plugins/card/base/org-info.tsx +++ b/web/app/components/plugins/card/base/org-info.tsx @@ -13,7 +13,7 @@ const OrgInfo = ({ className, orgName, packageName, packageNameClassName }: Prop {orgName && ( <> {orgName} @@ -23,7 +23,7 @@ const OrgInfo = ({ className, orgName, packageName, packageNameClassName }: Prop )} Date: Mon, 31 Aug 2026 05:51:30 +0000 Subject: [PATCH 11/21] test: migrate advanced chat generator sessions and ORM models to SQLite (#40585) --- .../apps/advanced_chat/test_app_generator.py | 619 +++++++++--------- 1 file changed, 311 insertions(+), 308 deletions(-) diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py index 2f69b1ed877..b99e0d23394 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py @@ -1,12 +1,15 @@ from __future__ import annotations +import json import logging from contextlib import contextmanager +from decimal import Decimal from types import SimpleNamespace from unittest.mock import MagicMock import pytest from pydantic import BaseModel, ValidationError +from sqlalchemy import Engine, event from sqlalchemy.orm import Session from constants import UUID_NIL @@ -21,20 +24,99 @@ from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom from core.ops.ops_trace_manager import TraceQueueManager from libs.datetime_utils import naive_utc_now -from models.enums import MessageStatus -from models.model import AppMode +from models.account import Account +from models.enums import ConversationFromSource, EndUserType, MessageStatus +from models.model import App, AppMode, Conversation, EndUser, Message +from models.workflow import Workflow, WorkflowType from tests.unit_tests.config_override import apply_config_overrides +def _make_app(*, app_id: str = "app", tenant_id: str = "tenant") -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name="Advanced Chat App", + mode=AppMode.ADVANCED_CHAT, + enable_site=False, + enable_api=False, + ) + + +def _make_workflow( + *, + workflow_id: str = "workflow-id", + tenant_id: str = "tenant", + app_id: str = "app", + features: dict[str, object] | None = None, +) -> Workflow: + return Workflow( + id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + type=WorkflowType.CHAT, + version=Workflow.VERSION_DRAFT, + graph="{}", + features=json.dumps(features or {}), + created_by="user", + ) + + +def _make_account(*, account_id: str = "user-id") -> Account: + account = Account(name="Advanced Chat User", email=f"{account_id}@example.com") + account.id = account_id + return account + + +def _make_end_user(*, end_user_id: str = "end-user-id", session_id: str = "session-id") -> EndUser: + return EndUser( + id=end_user_id, + tenant_id="tenant", + app_id="app", + type=EndUserType.BROWSER, + session_id=session_id, + ) + + +def _make_conversation(*, conversation_id: str = "conversation-id", app_id: str = "app") -> Conversation: + return Conversation( + id=conversation_id, + app_id=app_id, + mode=AppMode.ADVANCED_CHAT, + name="Advanced Chat Conversation", + inputs={}, + from_source=ConversationFromSource.API, + ) + + +def _make_message( + *, message_id: str = "message-id", conversation_id: str = "conversation-id", app_id: str = "app" +) -> Message: + return Message( + id=message_id, + app_id=app_id, + conversation_id=conversation_id, + inputs={}, + query="hello", + message={}, + answer="", + status=MessageStatus.NORMAL, + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + created_at=naive_utc_now(), + ) + + class TestAdvancedChatAppGeneratorValidation: def test_generate_requires_query(self, unbound_session: Session): generator = AdvancedChatAppGenerator() with pytest.raises(ValueError, match="query is required"): generator.generate( - app_model=SimpleNamespace(), - workflow=SimpleNamespace(), - user=SimpleNamespace(), + app_model=_make_app(), + workflow=_make_workflow(), + user=_make_account(), args={"inputs": {}}, invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", @@ -47,9 +129,9 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="query must be a string"): generator.generate( - app_model=SimpleNamespace(), - workflow=SimpleNamespace(), - user=SimpleNamespace(), + app_model=_make_app(), + workflow=_make_workflow(), + user=_make_account(), args={"inputs": {}, "query": 123}, invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", @@ -62,10 +144,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="node_id is required"): generator.single_iteration_generate( - app_model=SimpleNamespace(), - workflow=SimpleNamespace(), + app_model=_make_app(), + workflow=_make_workflow(), node_id="", - user=SimpleNamespace(), + user=_make_account(), args={"inputs": {}}, streaming=False, session=unbound_session, @@ -73,10 +155,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="inputs is required"): generator.single_iteration_generate( - app_model=SimpleNamespace(), - workflow=SimpleNamespace(), + app_model=_make_app(), + workflow=_make_workflow(), node_id="node", - user=SimpleNamespace(), + user=_make_account(), args={}, streaming=False, session=unbound_session, @@ -87,10 +169,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="node_id is required"): generator.single_loop_generate( - app_model=SimpleNamespace(), - workflow=SimpleNamespace(), + app_model=_make_app(), + workflow=_make_workflow(), node_id="", - user=SimpleNamespace(), + user=_make_account(), args=SimpleNamespace(inputs={}), streaming=False, session=unbound_session, @@ -98,10 +180,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="inputs is required"): generator.single_loop_generate( - app_model=SimpleNamespace(), - workflow=SimpleNamespace(), + app_model=_make_app(), + workflow=_make_workflow(), node_id="node", - user=SimpleNamespace(), + user=_make_account(), args=SimpleNamespace(inputs=None), streaming=False, session=unbound_session, @@ -120,11 +202,13 @@ class TestAdvancedChatAppGeneratorInternals: workflow_id="workflow-id", ) - def test_generate_loads_conversation_and_files(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): + def test_generate_loads_conversation_and_files( + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() app_config = self._build_app_config() - conversation = SimpleNamespace(id="conversation-id") + conversation = _make_conversation() built_files: list[object] = [] build_files_called = {"called": False} captured: dict[str, object] = {} @@ -157,10 +241,7 @@ class TestAdvancedChatAppGeneratorInternals: ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=lambda: SimpleNamespace(close=lambda: None)), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.sessionmaker", lambda **kwargs: SimpleNamespace() + SimpleNamespace(engine=sqlite_engine, session=unbound_session), ) monkeypatch.setattr(generator, "_prepare_user_inputs", lambda **kwargs: kwargs["user_inputs"]) @@ -187,8 +268,8 @@ class TestAdvancedChatAppGeneratorInternals: user.id = "user-id" result = generator.generate( - app_model=SimpleNamespace(id="app", tenant_id="tenant"), - workflow=SimpleNamespace(features_dict={}), + app_model=_make_app(), + workflow=_make_workflow(), user=user, args={ "query": "hello", @@ -238,11 +319,11 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.resume( - app_model=SimpleNamespace(id="app-id"), - workflow=SimpleNamespace(), - user=SimpleNamespace(id="end-user-id", session_id="session-id"), - conversation=SimpleNamespace(id="conversation-id"), - message=SimpleNamespace(id="message-id"), + app_model=_make_app(app_id="app-id"), + workflow=_make_workflow(), + user=_make_end_user(), + conversation=_make_conversation(), + message=_make_message(), session=unbound_session, application_generate_entity=application_generate_entity, workflow_execution_repository=SimpleNamespace(), @@ -257,7 +338,7 @@ class TestAdvancedChatAppGeneratorInternals: assert captured_graph_runtime_state is not None def test_single_iteration_generate_builds_debug_task( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine ): generator = AdvancedChatAppGenerator() app_config = self._build_app_config() @@ -265,7 +346,7 @@ class TestAdvancedChatAppGeneratorInternals: prefill_calls: list[object] = [] draft_sessions: list[object] = [] var_loader = SimpleNamespace(loader="draft") - workflow = SimpleNamespace(id="workflow-id") + workflow = _make_workflow() session = unbound_session monkeypatch.setattr( @@ -281,12 +362,9 @@ class TestAdvancedChatAppGeneratorInternals: lambda **kwargs: SimpleNamespace(repo="node"), ) monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.DraftVarLoader", lambda **kwargs: var_loader) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.sessionmaker", lambda **kwargs: SimpleNamespace() - ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=lambda: SimpleNamespace()), + SimpleNamespace(engine=sqlite_engine, session=unbound_session), ) class _DraftVarService: @@ -305,10 +383,10 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.single_iteration_generate( - app_model=SimpleNamespace(id="app", tenant_id="tenant"), + app_model=_make_app(), workflow=workflow, node_id="node-1", - user=SimpleNamespace(id="user-id"), + user=_make_account(), args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, streaming=False, session=session, @@ -322,14 +400,16 @@ class TestAdvancedChatAppGeneratorInternals: assert captured["application_generate_entity"].single_iteration_run.node_id == "node-1" assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" - def test_single_loop_generate_builds_debug_task(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): + def test_single_loop_generate_builds_debug_task( + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() app_config = self._build_app_config() captured: dict[str, object] = {} prefill_calls: list[object] = [] draft_sessions: list[object] = [] var_loader = SimpleNamespace(loader="draft") - workflow = SimpleNamespace(id="workflow-id") + workflow = _make_workflow() session = unbound_session monkeypatch.setattr( @@ -345,12 +425,9 @@ class TestAdvancedChatAppGeneratorInternals: lambda **kwargs: SimpleNamespace(repo="node"), ) monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.DraftVarLoader", lambda **kwargs: var_loader) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.sessionmaker", lambda **kwargs: SimpleNamespace() - ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=lambda: SimpleNamespace()), + SimpleNamespace(engine=sqlite_engine, session=unbound_session), ) class _DraftVarService: @@ -369,10 +446,10 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.single_loop_generate( - app_model=SimpleNamespace(id="app", tenant_id="tenant"), + app_model=_make_app(), workflow=workflow, node_id="node-2", - user=SimpleNamespace(id="user-id"), + user=_make_account(), args=SimpleNamespace(inputs={"foo": "bar"}, trace_session_id="session-1"), streaming=False, session=session, @@ -386,7 +463,9 @@ class TestAdvancedChatAppGeneratorInternals: assert captured["application_generate_entity"].single_loop_run.node_id == "node-2" assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" - def test_generate_internal_flow_initial_conversation_with_pause_layer(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_internal_flow_initial_conversation_with_pause_layer( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() generator._dialogue_count = 0 app_config = self._build_app_config() @@ -405,16 +484,19 @@ class TestAdvancedChatAppGeneratorInternals: workflow_run_id="run-id", ) - workflow = SimpleNamespace(id="wf-1", tenant_id="tenant", features={"feature": True}, features_dict={}) - conversation = SimpleNamespace(id="conv-1", mode=AppMode.ADVANCED_CHAT, override_model_configs=None) - message = SimpleNamespace( - id="msg-1", - query="hello", - created_at=naive_utc_now(), - status=MessageStatus.NORMAL, - answer="", - ) - db_session = SimpleNamespace(commit=MagicMock(), refresh=MagicMock(), close=MagicMock()) + app = _make_app() + workflow = _make_workflow(workflow_id="wf-1", features={"feature": True}) + conversation = _make_conversation(conversation_id="conv-1") + message = _make_message(message_id="msg-1", conversation_id=conversation.id) + sqlite_session.add_all([app, workflow, conversation, message]) + sqlite_session.commit() + commit_count = 0 + + def _record_commit(session: Session) -> None: + nonlocal commit_count + commit_count += 1 + + event.listen(sqlite_session, "after_commit", _record_commit) captured: dict[str, object] = {} thread_data: dict[str, object] = {} init_records = MagicMock(return_value=(conversation, message)) @@ -455,7 +537,8 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.threading.Thread", _Thread) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", SimpleNamespace(engine=object(), session=db_session) + "core.app.apps.advanced_chat.app_generator.db", + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory") monkeypatch.setattr( @@ -472,10 +555,10 @@ class TestAdvancedChatAppGeneratorInternals: response = generator._generate( workflow=workflow, - user=SimpleNamespace(id="user"), + user=_make_account(account_id="user"), invoke_from=InvokeFrom.WEB_APP, application_generate_entity=application_generate_entity, - session=db_session, + session=sqlite_session, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), conversation=None, @@ -490,17 +573,18 @@ class TestAdvancedChatAppGeneratorInternals: assert thread_data["join_timeout"] == 300 assert "pause-layer" in thread_data["kwargs"]["graph_engine_layers"] assert generator._dialogue_count == 3 - assert init_records.call_args.kwargs["session"] is db_session - get_thread_messages_length.assert_called_once_with(conversation.id, session=db_session) - db_session.commit.assert_called_once() - db_session.refresh.assert_called_once_with(conversation) - db_session.close.assert_called_once() + assert init_records.call_args.kwargs["session"] is sqlite_session + get_thread_messages_length.assert_called_once_with(conversation.id, session=sqlite_session) + assert commit_count == 1 + assert json.loads(conversation.override_model_configs) == {"feature": True} assert captured["draft_var_saver_factory"] == "draft-factory" assert isinstance(captured["workflow"], WorkflowSnapshot) assert isinstance(captured["conversation"], ConversationSnapshot) assert isinstance(captured["message"], MessageSnapshot) - def test_generate_internal_flow_with_existing_records_skips_init(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_internal_flow_with_existing_records_skips_init( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() generator._dialogue_count = 0 app_config = self._build_app_config() @@ -519,16 +603,19 @@ class TestAdvancedChatAppGeneratorInternals: workflow_run_id="run-id", ) - workflow = SimpleNamespace(id="wf-2", tenant_id="tenant", features={}, features_dict={}) - conversation = SimpleNamespace(id="conv-2", mode=AppMode.ADVANCED_CHAT, override_model_configs=None) - message = SimpleNamespace( - id="msg-2", - query="hello", - created_at=naive_utc_now(), - status=MessageStatus.NORMAL, - answer="", - ) - db_session = SimpleNamespace(close=MagicMock(), commit=MagicMock(), refresh=MagicMock()) + app = _make_app() + workflow = _make_workflow(workflow_id="wf-2") + conversation = _make_conversation(conversation_id="conv-2") + message = _make_message(message_id="msg-2", conversation_id=conversation.id) + sqlite_session.add_all([app, workflow, conversation, message]) + sqlite_session.commit() + commit_count = 0 + + def _record_commit(session: Session) -> None: + nonlocal commit_count + commit_count += 1 + + event.listen(sqlite_session, "after_commit", _record_commit) init_records = MagicMock() get_thread_messages_length = MagicMock(return_value=0) thread_data: dict[str, object] = {} @@ -564,7 +651,8 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.threading.Thread", _Thread) monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", SimpleNamespace(engine=object(), session=db_session) + "core.app.apps.advanced_chat.app_generator.db", + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory") monkeypatch.setattr( @@ -579,10 +667,10 @@ class TestAdvancedChatAppGeneratorInternals: response = generator._generate( workflow=workflow, - user=SimpleNamespace(id="user"), + user=_make_account(account_id="user"), invoke_from=InvokeFrom.WEB_APP, application_generate_entity=application_generate_entity, - session=db_session, + session=sqlite_session, workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), conversation=conversation, @@ -592,15 +680,15 @@ class TestAdvancedChatAppGeneratorInternals: assert response == {"raw": True} init_records.assert_not_called() - get_thread_messages_length.assert_called_once_with(conversation.id, session=db_session) + get_thread_messages_length.assert_called_once_with(conversation.id, session=sqlite_session) assert thread_data["started"] is True assert thread_data["joined"] is True assert thread_data["join_timeout"] == 300 - db_session.commit.assert_not_called() - db_session.refresh.assert_not_called() - db_session.close.assert_called_once() + assert commit_count == 0 - def test_generate_worker_raises_when_workflow_not_found(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_worker_raises_when_workflow_not_found( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() generator._dialogue_count = 1 app_config = self._build_app_config() @@ -619,8 +707,8 @@ class TestAdvancedChatAppGeneratorInternals: workflow_run_id="run-id", ) - generator._get_conversation = MagicMock(return_value=SimpleNamespace(id="conv")) - generator._get_message = MagicMock(return_value=SimpleNamespace(id="msg")) + generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) + generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) @contextmanager def _fake_context(*args, **kwargs): @@ -628,20 +716,9 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - class _Session: - def __init__(self, *args, **kwargs): - self.scalar = MagicMock(return_value=None) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) with pytest.raises(ValueError, match="Workflow not found"): @@ -659,7 +736,9 @@ class TestAdvancedChatAppGeneratorInternals: graph_runtime_state=None, ) - def test_generate_worker_raises_when_app_not_found_for_internal_call(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_worker_raises_when_app_not_found_for_internal_call( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() generator._dialogue_count = 1 app_config = self._build_app_config() @@ -678,8 +757,10 @@ class TestAdvancedChatAppGeneratorInternals: workflow_run_id="run-id", ) - generator._get_conversation = MagicMock(return_value=SimpleNamespace(id="conv")) - generator._get_message = MagicMock(return_value=SimpleNamespace(id="msg")) + generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) + generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) + sqlite_session.add(_make_workflow()) + sqlite_session.commit() @contextmanager def _fake_context(*args, **kwargs): @@ -687,25 +768,9 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - class _Session: - def __init__(self, *args, **kwargs): - self.scalar = MagicMock( - side_effect=[ - SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app"), - None, - ] - ) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) with pytest.raises(ValueError, match="App not found"): @@ -723,7 +788,9 @@ class TestAdvancedChatAppGeneratorInternals: graph_runtime_state=None, ) - def test_generate_worker_handles_stopped_error(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_worker_handles_stopped_error( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() generator._dialogue_count = 1 app_config = self._build_app_config() @@ -743,8 +810,8 @@ class TestAdvancedChatAppGeneratorInternals: ) queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=SimpleNamespace(id="conv")) - generator._get_message = MagicMock(return_value=SimpleNamespace(id="msg")) + generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) + generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) @contextmanager def _fake_context(*args, **kwargs): @@ -752,22 +819,8 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - workflow = SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app") - - class _Session: - def __init__(self, *args, **kwargs): - self.scalar = MagicMock( - side_effect=[ - workflow, - SimpleNamespace(id="app"), - ] - ) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False + sqlite_session.add_all([_make_app(), _make_workflow()]) + sqlite_session.commit() class _Runner: def __init__(self, **kwargs): @@ -776,13 +829,12 @@ class TestAdvancedChatAppGeneratorInternals: def run(self): raise GenerateTaskStoppedError() - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session) monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", _Runner) restore_workflow_run_graph = MagicMock() monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) generator._generate_worker( @@ -800,10 +852,12 @@ class TestAdvancedChatAppGeneratorInternals: ) queue_manager.publish_error.assert_not_called() - assert restore_workflow_run_graph.call_args.kwargs["workflow"] is workflow + assert restore_workflow_run_graph.call_args.kwargs["workflow"].id == "workflow-id" assert restore_workflow_run_graph.call_args.kwargs["workflow_run_id"] == "run-id" - def test_generate_worker_handles_validation_error(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_worker_handles_validation_error( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() generator._dialogue_count = 1 app_config = self._build_app_config() @@ -833,8 +887,10 @@ class TestAdvancedChatAppGeneratorInternals: raise AssertionError("validation error should be created") queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=SimpleNamespace(id="conv")) - generator._get_message = MagicMock(return_value=SimpleNamespace(id="msg")) + generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) + generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) + sqlite_session.add_all([_make_app(), _make_workflow()]) + sqlite_session.commit() @contextmanager def _fake_context(*args, **kwargs): @@ -842,21 +898,6 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - class _Session: - def __init__(self, *args, **kwargs): - self.scalar = MagicMock( - side_effect=[ - SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app"), - SimpleNamespace(id="app"), - ] - ) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - class _Runner: def __init__(self, **kwargs): _ = kwargs @@ -864,11 +905,10 @@ class TestAdvancedChatAppGeneratorInternals: def run(self): raise validation_error - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session) monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", _Runner) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) generator._generate_worker( @@ -887,8 +927,12 @@ class TestAdvancedChatAppGeneratorInternals: queue_manager.publish_error.assert_called_once() - def test_generate_worker_handles_value_and_unknown_errors(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_worker_handles_value_and_unknown_errors( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): app_config = self._build_app_config() + sqlite_session.add_all([_make_app(), _make_workflow()]) + sqlite_session.commit() @contextmanager def _fake_context(*args, **kwargs): @@ -922,26 +966,10 @@ class TestAdvancedChatAppGeneratorInternals: ) queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=SimpleNamespace(id="conv")) - generator._get_message = MagicMock(return_value=SimpleNamespace(id="msg")) - - class _Session: - def __init__(self, *args, **kwargs): - self.scalar = MagicMock( - side_effect=[ - SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app"), - SimpleNamespace(id="app"), - ] - ) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False + generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) + generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", _make_runner(raised_error), @@ -949,7 +977,7 @@ class TestAdvancedChatAppGeneratorInternals: apply_config_overrides(monkeypatch, DEBUG=True) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) generator._generate_worker( @@ -1019,7 +1047,7 @@ class TestAdvancedChatAppGeneratorInternals: status=MessageStatus.NORMAL, answer="", ), - user=SimpleNamespace(), + user=_make_account(), draft_var_saver_factory=lambda **kwargs: None, stream=False, ) @@ -1067,14 +1095,16 @@ class TestAdvancedChatAppGeneratorInternals: status=MessageStatus.NORMAL, answer="", ), - user=SimpleNamespace(), + user=_make_account(), draft_var_saver_factory=lambda **kwargs: None, stream=False, ) assert "Failed to process generate task pipeline, conversation_id: conv" in caplog.messages - def test_generate_worker_handles_invoke_auth_error(self, monkeypatch: pytest.MonkeyPatch): + def test_generate_worker_handles_invoke_auth_error( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + ): generator = AdvancedChatAppGenerator() generator._dialogue_count = 1 @@ -1102,8 +1132,10 @@ class TestAdvancedChatAppGeneratorInternals: queue_manager = MagicMock() - generator._get_conversation = MagicMock(return_value=SimpleNamespace(id="conv", mode=AppMode.ADVANCED_CHAT)) - generator._get_message = MagicMock(return_value=SimpleNamespace(id="msg")) + generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) + generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) + sqlite_session.add_all([_make_app(), _make_workflow(), _make_end_user()]) + sqlite_session.commit() class _Runner: def __init__(self, **kwargs) -> None: @@ -1122,26 +1154,9 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - class _Session: - def __init__(self, *args, **kwargs): - self.scalar = MagicMock( - side_effect=[ - SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app"), - SimpleNamespace(id="end-user-id", session_id="session-id"), - SimpleNamespace(id="app"), - ] - ) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), ) generator._generate_worker( @@ -1160,88 +1175,8 @@ class TestAdvancedChatAppGeneratorInternals: assert queue_manager.publish_error.called - def test_generate_debugger_enables_retrieve_source(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): - generator = AdvancedChatAppGenerator() - - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.ADVANCED_CHAT, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", - lambda app_model, workflow: app_config, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", - lambda features_dict, is_vision=False: None, - ) - DummyTraceQueueManager = type( - "_DummyTraceQueueManager", - (TraceQueueManager,), - { - "__init__": lambda self, app_id=None, user_id=None: ( - setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) - ) - }, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", - DummyTraceQueueManager, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.sessionmaker", - lambda **kwargs: SimpleNamespace(), - ) - - captured = {} - - def _fake_generate(**kwargs): - captured.update(kwargs) - return {"ok": True} - - monkeypatch.setattr(generator, "_generate", _fake_generate) - - app_model = SimpleNamespace(id="app", tenant_id="tenant") - workflow = SimpleNamespace(features_dict={}) - from models import Account - - user = Account(name="Tester", email="tester@example.com") - user.id = "user" - - result = generator.generate( - app_model=app_model, - workflow=workflow, - user=user, - args={"query": "hello\x00", "inputs": {}}, - invoke_from=InvokeFrom.DEBUGGER, - workflow_run_id="run-id", - streaming=False, - session=unbound_session, - ) - - assert result == {"ok": True} - assert app_config.additional_features.show_retrieve_source is True - assert captured["application_generate_entity"].query == "hello" - - def test_generate_service_api_sets_parent_message_id( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + def test_generate_debugger_enables_retrieve_source( + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine ): generator = AdvancedChatAppGenerator() @@ -1285,11 +1220,7 @@ class TestAdvancedChatAppGeneratorInternals: ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.sessionmaker", - lambda **kwargs: SimpleNamespace(), + SimpleNamespace(engine=sqlite_engine, session=unbound_session), ) captured = {} @@ -1300,12 +1231,84 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) - app_model = SimpleNamespace(id="app", tenant_id="tenant") - workflow = SimpleNamespace(features_dict={}) - from models.model import EndUser + app_model = _make_app() + workflow = _make_workflow() + user = _make_account(account_id="user") - user = EndUser(tenant_id="tenant", type="session", name="tester", session_id="session") - user.id = "end-user" + result = generator.generate( + app_model=app_model, + workflow=workflow, + user=user, + args={"query": "hello\x00", "inputs": {}}, + invoke_from=InvokeFrom.DEBUGGER, + workflow_run_id="run-id", + streaming=False, + session=unbound_session, + ) + + assert result == {"ok": True} + assert app_config.additional_features.show_retrieve_source is True + assert captured["application_generate_entity"].query == "hello" + + def test_generate_service_api_sets_parent_message_id( + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + ): + generator = AdvancedChatAppGenerator() + + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.ADVANCED_CHAT, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + lambda app_model, workflow: app_config, + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", + lambda features_dict, is_vision=False: None, + ) + DummyTraceQueueManager = type( + "_DummyTraceQueueManager", + (TraceQueueManager,), + { + "__init__": lambda self, app_id=None, user_id=None: ( + setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) + ) + }, + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.TraceQueueManager", + DummyTraceQueueManager, + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", + lambda **kwargs: SimpleNamespace(), + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + lambda **kwargs: SimpleNamespace(), + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.db", + SimpleNamespace(engine=sqlite_engine, session=unbound_session), + ) + + captured = {} + + def _fake_generate(**kwargs): + captured.update(kwargs) + return {"ok": True} + + monkeypatch.setattr(generator, "_generate", _fake_generate) + + app_model = _make_app() + workflow = _make_workflow() + user = _make_end_user(end_user_id="end-user", session_id="session") generator.generate( app_model=app_model, @@ -1376,11 +1379,11 @@ class TestAdvancedChatAppGeneratorResume: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.resume( - app_model=SimpleNamespace(id="app-id"), - workflow=SimpleNamespace(), - user=SimpleNamespace(id="end-user-id", session_id="session-id"), - conversation=SimpleNamespace(id="conversation-id"), - message=SimpleNamespace(id="message-id"), + app_model=_make_app(app_id="app-id"), + workflow=_make_workflow(), + user=_make_end_user(), + conversation=_make_conversation(), + message=_make_message(), session=unbound_session, application_generate_entity=application_generate_entity, workflow_execution_repository=SimpleNamespace(), @@ -1424,11 +1427,11 @@ class TestAdvancedChatAppGeneratorResume: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.resume( - app_model=SimpleNamespace(id="app-id"), - workflow=SimpleNamespace(), - user=SimpleNamespace(id="end-user-id", session_id="session-id"), - conversation=SimpleNamespace(id="conversation-id"), - message=SimpleNamespace(id="message-id"), + app_model=_make_app(app_id="app-id"), + workflow=_make_workflow(), + user=_make_end_user(), + conversation=_make_conversation(), + message=_make_message(), session=unbound_session, application_generate_entity=application_generate_entity, workflow_execution_repository=SimpleNamespace(), From 9398482b53469be5ed60acbf0a3fc58dde23a687 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 31 Aug 2026 05:55:12 +0000 Subject: [PATCH 12/21] test: migrate console dataset sessions and ORM models to SQLite (#40521) --- .../console/datasets/test_datasets.py | 327 ++++++++++-------- 1 file changed, 174 insertions(+), 153 deletions(-) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py index ada7de0ce20..b87cdc3ecaa 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py @@ -8,6 +8,7 @@ from unittest.mock import ANY, MagicMock, PropertyMock, call, patch import pytest from flask import Flask +from sqlalchemy.orm import Session from werkzeug.exceptions import BadRequest, Forbidden, NotFound import services @@ -44,7 +45,7 @@ from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.retrieval.retrieval_methods import RetrievalMethod from extensions.storage.storage_type import StorageType from models.account import Account, TenantAccountRole -from models.dataset import Dataset, DatasetQuery, Document +from models.dataset import AppDatasetJoin, Dataset, DatasetPermission, DatasetQuery, Document, DocumentSegment from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import ApiToken, App, AppMode, IconType, UploadFile from services.dataset_ref_service import DatasetRef @@ -170,7 +171,29 @@ def make_document_status(**overrides) -> Document: return Document(**base) -class TestDatasetList: +def make_document_segment(*, position: int, completed: bool) -> DocumentSegment: + return DocumentSegment( + tenant_id="tenant-1", + dataset_id="dataset-1", + document_id="doc-1", + position=position, + content=f"segment {position}", + word_count=2, + tokens=2, + created_by="account-1", + completed_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC) if completed else None, + ) + + +class _UsesSQLiteSession: + session: Session + + @pytest.fixture(autouse=True) + def _inject_sqlite_session(self, sqlite_session: Session) -> None: + self.session = sqlite_session + + +class TestDatasetList(_UsesSQLiteSession): def _mock_user(self): user = make_account() return user @@ -185,7 +208,7 @@ class TestDatasetList: patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - resp, status = method(api, MagicMock(), "tenant-1", current_user) + resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 assert resp["total"] == 1 assert resp["data"][0]["embedding_available"] is True @@ -201,7 +224,7 @@ class TestDatasetList: method = unwrap(api.get) current_user = self._mock_user() dataset = make_dataset() - session = MagicMock() + session = self.session with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=([dataset], 1)), @@ -222,7 +245,7 @@ class TestDatasetList: patch.object(DatasetService, "get_datasets_by_ids", return_value=(datasets, 2)) as by_ids_mock, patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - resp, status = method(api, MagicMock(), "tenant-1", current_user) + resp, status = method(api, self.session, "tenant-1", current_user) by_ids_mock.assert_called_once() assert status == 200 assert resp["total"] == 2 @@ -251,7 +274,7 @@ class TestDatasetList: return_value=permissions, ) as get_permissions, ): - resp, status = method(api, MagicMock(), "tenant-1", current_user) + resp, status = method(api, self.session, "tenant-1", current_user) get_permissions.assert_called_once_with("tenant-1", current_user.id, session=ANY) assert status == 200 assert resp["data"][0]["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] @@ -281,7 +304,7 @@ class TestDatasetList: ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - method(api, MagicMock(), "tenant-1", current_user) + method(api, self.session, "tenant-1", current_user) assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == [] assert get_datasets.call_args.kwargs["include_own_datasets"] is False @@ -308,7 +331,7 @@ class TestDatasetList: ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - method(api, MagicMock(), "tenant-1", current_user) + method(api, self.session, "tenant-1", current_user) assert get_datasets.call_args.kwargs["accessible_dataset_ids"] is None def test_get_restricted_whitelist_overrides_default_read_permission( @@ -374,7 +397,7 @@ class TestDatasetList: ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - method(api, MagicMock(), "tenant-1", current_user) + method(api, self.session, "tenant-1", current_user) assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == [ "dataset-whitelist-only", ] @@ -399,9 +422,9 @@ class TestDatasetList: ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - method(api, MagicMock(), "tenant-1", current_user) + method(api, self.session, "tenant-1", current_user) session = get_datasets_by_ids.call_args.kwargs["session"] - assert isinstance(session, MagicMock) + assert session is self.session assert get_datasets_by_ids.call_args.args == (["dataset-1"], "tenant-1") assert get_datasets_by_ids.call_args.kwargs == { "user": current_user, @@ -420,7 +443,7 @@ class TestDatasetList: patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - resp, status = method(api, MagicMock(), "tenant-1", current_user) + resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 def test_get_allows_legacy_weighted_score_without_weight_type(self, app: Flask): @@ -453,7 +476,7 @@ class TestDatasetList: patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - resp, status = method(api, MagicMock(), "tenant-1", current_user) + resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 assert resp["data"][0]["retrieval_model_dict"]["weights"]["weight_type"] is None @@ -467,7 +490,7 @@ class TestDatasetList: patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): - resp, status = method(api, MagicMock(), "tenant-1", current_user) + resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 retrieval_model = resp["data"][0]["retrieval_model_dict"] assert retrieval_model["search_method"] == "semantic_search" @@ -491,7 +514,7 @@ class TestDatasetList: patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=config), ): - resp, status = method(api, MagicMock(), "tenant-1", current_user) + resp, status = method(api, self.session, "tenant-1", current_user) assert resp["data"][0]["embedding_available"] is False def test_partial_members_permission(self, app: Flask): @@ -499,8 +522,9 @@ class TestDatasetList: method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset(permission="partial_members")] - session = MagicMock() - session.execute.return_value.all.return_value = [("ds-1", "u1")] + session = self.session + session.add(DatasetPermission(dataset_id="ds-1", account_id="u1", tenant_id="tenant-1")) + session.flush() with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), @@ -510,7 +534,7 @@ class TestDatasetList: assert resp["data"][0]["partial_member_list"] == ["u1"] -class TestDatasetListApiPost: +class TestDatasetListApiPost(_UsesSQLiteSession): def test_post_success(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) @@ -522,7 +546,7 @@ class TestDatasetListApiPost: patch.object(type(console_ns), "payload", payload), patch.object(DatasetService, "create_empty_dataset", return_value=dataset), ): - _, status = method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user) + _, status = method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user) assert status == 201 def test_post_forbidden(self, app: Flask): @@ -532,7 +556,7 @@ class TestDatasetListApiPost: user = make_account(TenantAccountRole.NORMAL) with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(Forbidden): - method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user) + method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user) def test_post_duplicate_name(self, app: Flask): api = DatasetListApi() @@ -547,14 +571,14 @@ class TestDatasetListApiPost: ), ): with pytest.raises(DatasetNameDuplicateError): - method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user) + method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user) def test_post_invalid_payload_missing_name(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) with app.test_request_context("/datasets", json={}), patch.object(type(console_ns), "payload", {}): with pytest.raises(ValueError): - method(api, DatasetCreatePayload(), MagicMock(), "tenant-1", make_account()) + method(api, DatasetCreatePayload(), self.session, "tenant-1", make_account()) def test_post_invalid_indexing_technique(self, app: Flask): api = DatasetListApi() @@ -562,7 +586,7 @@ class TestDatasetListApiPost: payload = {"name": "bad", "indexing_technique": "invalid-tech"} with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(ValueError, match="Invalid indexing technique"): - method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account()) + method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account()) def test_post_invalid_provider(self, app: Flask): api = DatasetListApi() @@ -570,10 +594,10 @@ class TestDatasetListApiPost: payload = {"name": "bad", "provider": "unknown"} with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(ValueError, match="Invalid provider"): - method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account()) + method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account()) -class TestDatasetApiGet: +class TestDatasetApiGet(_UsesSQLiteSession): def test_get_success_basic(self, app: Flask): api = DatasetApi() method = unwrap(api.get) @@ -588,7 +612,7 @@ class TestDatasetApiGet: patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, status = method(api, MagicMock(), tenant_id, user, dataset_id) + data, status = method(api, self.session, tenant_id, user, dataset_id) assert status == 200 assert data["embedding_available"] is True @@ -597,7 +621,7 @@ class TestDatasetApiGet: api = DatasetApi() method = unwrap(api.get) dataset_id = "123e4567-e89b-12d3-a456-426614174000" - user = MagicMock(id="account-1") + user = make_account() tenant_id = "tenant-1" dataset = make_dataset(id=dataset_id) with ( @@ -619,7 +643,7 @@ class TestDatasetApiGet: patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, status = method(api, MagicMock(), tenant_id, user, dataset_id) + data, status = method(api, self.session, tenant_id, user, dataset_id) get_permissions.assert_called_once_with(tenant_id, user.id, dataset_id=dataset_id, session=ANY) assert status == 200 assert data["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] @@ -636,7 +660,7 @@ class TestDatasetApiGet: patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, status = method(api, MagicMock(), "tenant", make_account(), dataset_id) + data, status = method(api, self.session, "tenant", make_account(), dataset_id) assert status == 200 assert data["external_retrieval_model"] == {"top_k": 2, "score_threshold": 0.0, "score_threshold_enabled": None} @@ -649,7 +673,7 @@ class TestDatasetApiGet: patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): - method(api, MagicMock(), "tenant", make_account(), dataset_id) + method(api, self.session, "tenant", make_account(), dataset_id) def test_get_permission_denied(self, app: Flask): api = DatasetApi() @@ -666,7 +690,7 @@ class TestDatasetApiGet: ), ): with pytest.raises(Forbidden, match="no access"): - method(api, MagicMock(), "tenant", make_account(), dataset_id) + method(api, self.session, "tenant", make_account(), dataset_id) def test_get_high_quality_embedding_unavailable(self, app: Flask): api = DatasetApi() @@ -687,7 +711,7 @@ class TestDatasetApiGet: patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, _ = method(api, MagicMock(), tenant_id, user, dataset_id) + data, _ = method(api, self.session, tenant_id, user, dataset_id) assert data["embedding_available"] is False def test_get_partial_members_permission(self, app: Flask): @@ -704,11 +728,11 @@ class TestDatasetApiGet: patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, _ = method(api, MagicMock(), "tenant", make_account(), dataset_id) + data, _ = method(api, self.session, "tenant", make_account(), dataset_id) assert data["partial_member_list"] == partial_members -class TestDatasetApiPatch: +class TestDatasetApiPatch(_UsesSQLiteSession): def test_patch_success_basic(self, app: Flask): api = DatasetApi() method = unwrap(api.patch) @@ -725,7 +749,7 @@ class TestDatasetApiPatch: patch.object(DatasetService, "update_dataset", return_value=dataset), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), ): - result, status = method(api, DatasetUpdatePayload(), MagicMock(), tenant_id, user, dataset_id) + result, status = method(api, DatasetUpdatePayload(), self.session, tenant_id, user, dataset_id) assert status == 200 assert result["partial_member_list"] == [] @@ -737,7 +761,7 @@ class TestDatasetApiPatch: patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): - method(api, DatasetUpdatePayload(), MagicMock(), "tenant-1", make_account(), "missing") + method(api, DatasetUpdatePayload(), self.session, "tenant-1", make_account(), "missing") def test_patch_permission_denied(self, app: Flask): api = DatasetApi() @@ -752,7 +776,7 @@ class TestDatasetApiPatch: patch.object(DatasetPermissionService, "check_permission", side_effect=Forbidden("no permission")), ): with pytest.raises(Forbidden): - method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id) + method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) def test_patch_partial_members_update(self, app: Flask): api = DatasetApi() @@ -769,7 +793,7 @@ class TestDatasetApiPatch: patch.object(DatasetPermissionService, "update_partial_member_list", return_value=None), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["u1", "u2"]), ): - result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id) + result, _ = method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) assert result["partial_member_list"] == ["u1", "u2"] def test_patch_clear_partial_members(self, app: Flask): @@ -787,11 +811,11 @@ class TestDatasetApiPatch: patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), ): - result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id) + result, _ = method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) assert result["partial_member_list"] == [] -class TestDatasetApiDelete: +class TestDatasetApiDelete(_UsesSQLiteSession): def test_delete_success(self, app: Flask): api = DatasetApi() method = unwrap(api.delete) @@ -802,7 +826,7 @@ class TestDatasetApiDelete: patch.object(DatasetService, "delete_dataset", return_value=True), patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None), ): - result, status = method(api, MagicMock(), user, dataset_id) + result, status = method(api, self.session, user, dataset_id) assert status == 204 assert result == "" @@ -813,7 +837,7 @@ class TestDatasetApiDelete: user = make_account(TenantAccountRole.NORMAL) with app.test_request_context(f"/datasets/{dataset_id}"): with pytest.raises(Forbidden): - method(api, MagicMock(), user, dataset_id) + method(api, self.session, user, dataset_id) def test_delete_dataset_not_found(self, app: Flask): api = DatasetApi() @@ -825,7 +849,7 @@ class TestDatasetApiDelete: patch.object(DatasetService, "delete_dataset", return_value=False), ): with pytest.raises(NotFound, match="Dataset not found"): - method(api, MagicMock(), user, dataset_id) + method(api, self.session, user, dataset_id) def test_delete_dataset_in_use(self, app: Flask): api = DatasetApi() @@ -837,10 +861,10 @@ class TestDatasetApiDelete: patch.object(DatasetService, "delete_dataset", side_effect=services.errors.dataset.DatasetInUseError()), ): with pytest.raises(DatasetInUseError): - method(api, MagicMock(), user, dataset_id) + method(api, self.session, user, dataset_id) -class TestDatasetUseCheckApi: +class TestDatasetUseCheckApi(_UsesSQLiteSession): @pytest.mark.parametrize("is_using", [True, False]) def test_get_use_check(self, app: Flask, is_using: bool): api = DatasetUseCheckApi() @@ -848,7 +872,7 @@ class TestDatasetUseCheckApi: dataset_id = "dataset-id" dataset = make_dataset(id=dataset_id) current_user = make_account() - session = MagicMock() + session = self.session with ( app.test_request_context(f"/datasets/{dataset_id}/use-check"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, @@ -867,7 +891,7 @@ class TestDatasetUseCheckApi: api = DatasetUseCheckApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-id") - session = MagicMock() + session = self.session with ( app.test_request_context("/datasets/dataset-id/use-check"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), @@ -884,11 +908,11 @@ class TestDatasetUseCheckApi: "api_cls", [DatasetUseCheckApi, DatasetIndexingStatusApi, DatasetErrorDocs, DatasetAutoDisableLogApi], ) -def test_dataset_scoped_read_permission_denied(app: Flask, api_cls): +def test_dataset_scoped_read_permission_denied(app: Flask, api_cls, sqlite_session: Session): api = api_cls() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") - session = MagicMock() + session = sqlite_session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), @@ -902,7 +926,7 @@ def test_dataset_scoped_read_permission_denied(app: Flask, api_cls): method(api, session, "tenant-1", make_account(), "dataset-1") -class TestDatasetQueryApi: +class TestDatasetQueryApi(_UsesSQLiteSession): def _query_record(self, index: int = 1) -> DatasetQuery: query = DatasetQuery( dataset_id="dataset-id", @@ -929,7 +953,7 @@ class TestDatasetQueryApi: patch.object(DatasetService, "check_dataset_permission", return_value=None), patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 2)), ): - response, status = method(api, MagicMock(), current_user, dataset_id) + response, status = method(api, self.session, current_user, dataset_id) assert status == 200 assert response["total"] == 2 assert response["page"] == 1 @@ -952,24 +976,30 @@ class TestDatasetQueryApi: dataset = make_dataset(id="dataset-id") query = self._query_record() query.content = json.dumps([{"content_type": "image_query", "content": "file-1"}]) - upload_file = SimpleNamespace( - id="file-1", + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key="image.png", name="image.png", size=10, extension="png", mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC), + used=False, ) - session = MagicMock() - session.scalar.return_value = upload_file + upload_file.id = "file-1" + session = self.session + session.add(upload_file) + session.flush() with ( app.test_request_context("/datasets/queries"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch.object(DatasetService, "get_dataset_queries", return_value=([query], 1)), - patch("models.dataset.db") as db_mock, patch("models.dataset.sign_upload_file_preview_url", return_value="signed-url"), ): - db_mock.session.scalar.return_value = upload_file response, status = method(api, session, make_account(), "dataset-id") assert status == 200 @@ -987,8 +1017,7 @@ class TestDatasetQueryApi: }, } ] - session.scalar.assert_called_once() - db_mock.session.scalar.assert_not_called() + assert session.get(UploadFile, "file-1") is upload_file def test_get_queries_dataset_not_found(self, app: Flask): api = DatasetQueryApi() @@ -1000,7 +1029,7 @@ class TestDatasetQueryApi: patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): - method(api, MagicMock(), current_user, dataset_id) + method(api, self.session, current_user, dataset_id) def test_get_queries_permission_denied(self, app: Flask): api = DatasetQueryApi() @@ -1018,7 +1047,7 @@ class TestDatasetQueryApi: ), ): with pytest.raises(Forbidden): - method(api, MagicMock(), current_user, dataset_id) + method(api, self.session, current_user, dataset_id) def test_get_queries_pagination_has_more(self, app: Flask): api = DatasetQueryApi() @@ -1033,13 +1062,13 @@ class TestDatasetQueryApi: patch.object(DatasetService, "check_dataset_permission", return_value=None), patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 40)), ): - response, status = method(api, MagicMock(), current_user, dataset_id) + response, status = method(api, self.session, current_user, dataset_id) assert status == 200 assert response["has_more"] is True assert len(response["data"]) == 20 -class TestDatasetIndexingEstimateApi: +class TestDatasetIndexingEstimateApi(_UsesSQLiteSession): def _upload_file(self, *, tenant_id: str = "tenant-1", file_id: str = "file-1") -> UploadFile: upload_file = UploadFile( tenant_id=tenant_id, @@ -1072,8 +1101,9 @@ class TestDatasetIndexingEstimateApi: method = unwrap(api.post) payload = self._base_payload() mock_file = self._upload_file() - session = MagicMock() - session.scalars.return_value.all.return_value = [mock_file] + session = self.session + session.add(mock_file) + session.flush() mock_response = IndexingEstimate(total_segments=100, preview=[]) @@ -1102,8 +1132,7 @@ class TestDatasetIndexingEstimateApi: api = DatasetIndexingEstimateApi() method = unwrap(api.post) payload = self._base_payload() - session = MagicMock() - session.scalars.return_value.all.return_value = None + session = self.session with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), @@ -1122,8 +1151,9 @@ class TestDatasetIndexingEstimateApi: method = unwrap(api.post) mock_file = self._upload_file() payload = self._base_payload() - session = MagicMock() - session.scalars.return_value.all.return_value = [mock_file] + session = self.session + session.add(mock_file) + session.flush() with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), @@ -1146,8 +1176,9 @@ class TestDatasetIndexingEstimateApi: method = unwrap(api.post) mock_file = self._upload_file() payload = self._base_payload() - session = MagicMock() - session.scalars.return_value.all.return_value = [mock_file] + session = self.session + session.add(mock_file) + session.flush() with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), @@ -1170,8 +1201,9 @@ class TestDatasetIndexingEstimateApi: method = unwrap(api.post) mock_file = self._upload_file() payload = self._base_payload() - session = MagicMock() - session.scalars.return_value.all.return_value = [mock_file] + session = self.session + session.add(mock_file) + session.flush() with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), @@ -1189,16 +1221,16 @@ class TestDatasetIndexingEstimateApi: ) -class TestDatasetRelatedAppListApi: +class TestDatasetRelatedAppListApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetRelatedAppListApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") app1 = make_related_app(id="app-1", name="App 1") app2 = make_related_app(id="app-2", name="App 2") - join1 = MagicMock(app_id="app-1") - join2 = MagicMock(app_id="app-2") - session = MagicMock() + join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1") + join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1") + session = self.session with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), @@ -1251,7 +1283,7 @@ class TestDatasetRelatedAppListApi: patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=None), ): with pytest.raises(NotFound): - method(api, MagicMock(), make_account(), "dataset-1") + method(api, self.session, make_account(), "dataset-1") def test_get_permission_denied(self, app: Flask): api = DatasetRelatedAppListApi() @@ -1266,16 +1298,16 @@ class TestDatasetRelatedAppListApi: ), ): with pytest.raises(Forbidden): - method(api, MagicMock(), make_account(), "dataset-1") + method(api, self.session, make_account(), "dataset-1") def test_get_filters_none_apps(self, app: Flask): api = DatasetRelatedAppListApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") app1 = make_related_app() - join1 = MagicMock(app_id="app-1") - join2 = MagicMock(app_id="app-2") - session = MagicMock() + join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1") + join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1") + session = self.session with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), @@ -1303,26 +1335,17 @@ class TestDatasetRelatedAppListApi: ] -class TestDatasetIndexingStatusApi: +class TestDatasetIndexingStatusApi(_UsesSQLiteSession): def test_get_success_with_documents(self, app: Flask): api = DatasetIndexingStatusApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") current_user = make_account() - document = MagicMock() - document.id = "doc-1" - document.indexing_status = "completed" - document.processing_started_at = None - document.parsing_completed_at = None - document.cleaning_completed_at = None - document.splitting_completed_at = None - document.completed_at = None - document.paused_at = None - document.error = None - document.stopped_at = None - session = MagicMock() - session.scalars.return_value.all.return_value = [document] - session.scalar.return_value = 3 + document = make_document_status() + session = self.session + session.add(document) + session.add_all([make_document_segment(position=position, completed=True) for position in range(1, 4)]) + session.flush() with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, @@ -1337,16 +1360,13 @@ class TestDatasetIndexingStatusApi: assert item["total_segments"] == 3 get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) check_permission.assert_called_once_with(dataset, current_user, session) - assert {"dataset-1", "tenant-1"} <= set(session.scalars.call_args.args[0].compile().params.values()) - for segment_count_call in session.scalar.call_args_list: - assert {"dataset-1", "tenant-1", "doc-1"} <= set(segment_count_call.args[0].compile().params.values()) + assert session.get(Document, "doc-1") is document def test_get_success_no_documents(self, app: Flask): api = DatasetIndexingStatusApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") - session = MagicMock() - session.scalars.return_value.all.return_value = [] + session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), @@ -1360,20 +1380,11 @@ class TestDatasetIndexingStatusApi: api = DatasetIndexingStatusApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") - document = MagicMock() - document.id = "doc-1" - document.indexing_status = "indexing" - document.processing_started_at = None - document.parsing_completed_at = None - document.cleaning_completed_at = None - document.splitting_completed_at = None - document.completed_at = None - document.paused_at = None - document.error = None - document.stopped_at = None - session = MagicMock() - session.scalars.return_value.all.return_value = [document] - session.scalar.side_effect = [2, 5] + document = make_document_status(indexing_status=IndexingStatus.INDEXING) + session = self.session + session.add(document) + session.add_all([make_document_segment(position=position, completed=position <= 2) for position in range(1, 6)]) + session.flush() with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), @@ -1386,7 +1397,7 @@ class TestDatasetIndexingStatusApi: assert item["total_segments"] == 5 -class TestDatasetApiKeyApi: +class TestDatasetApiKeyApi(_UsesSQLiteSession): def test_get_api_keys_success(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.get) @@ -1404,8 +1415,11 @@ class TestDatasetApiKeyApi: last_used_at=None, created_at=None, ) - session = MagicMock() - session.scalars.return_value.all.return_value = [mock_key_1, mock_key_2] + session = self.session + mock_key_1.tenant_id = "tenant-1" + mock_key_2.tenant_id = "tenant-1" + session.add_all([mock_key_1, mock_key_2]) + session.flush() with app.test_request_context("/"): response = method(api, session, "tenant-1") assert "data" in response @@ -1418,30 +1432,31 @@ class TestDatasetApiKeyApi: def test_post_create_api_key_success(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) - mock_token = MagicMock() - mock_token.id = "new-key-id" - mock_token.last_used_at = None - mock_token.created_at = datetime.datetime(2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC) - mock_api_token_cls = MagicMock() - mock_api_token_cls.return_value = mock_token - mock_api_token_cls.generate_api_key.return_value = "dataset-abc123" - session = MagicMock() - session.scalar.return_value = 3 - with app.test_request_context("/"), patch("controllers.console.datasets.datasets.ApiToken", mock_api_token_cls): + session = self.session + with ( + app.test_request_context("/"), + patch.object(ApiToken, "generate_api_key", return_value="dataset-abc123") as generate_api_key, + ): response, status = method(api, session, "tenant-1") assert status == 200 assert isinstance(response, dict) - assert response["id"] == "new-key-id" assert response["token"] == "dataset-abc123" assert response["type"] == "dataset" assert response["created_at"] is not None - mock_api_token_cls.generate_api_key.assert_called_once_with("dataset-", 24, session=session) + generate_api_key.assert_called_once_with("dataset-", 24, session=session) + assert session.get(ApiToken, response["id"]).token == "dataset-abc123" def test_post_exceed_max_keys(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) - session = MagicMock() - session.scalar.return_value = 10 + session = self.session + session.add_all( + [ + ApiToken(id=f"key-{index}", tenant_id="tenant-1", type="dataset", token=f"ds-{index}") + for index in range(10) + ] + ) + session.flush() with app.test_request_context("/"): with pytest.raises(BadRequest) as exc_info: method(api, session, "tenant-1") @@ -1452,36 +1467,42 @@ class TestDatasetApiKeyApi: } -class TestDatasetApiDeleteApi: +class TestDatasetApiDeleteApi(_UsesSQLiteSession): def test_delete_success(self, app: Flask): api = DatasetApiDeleteApi() method = unwrap(api.delete) - mock_key = MagicMock() - session = MagicMock() - session.scalar.return_value = mock_key - with app.test_request_context("/"): + session = self.session + key = ApiToken(id="api-key-id", tenant_id="tenant-1", type="dataset", token="dataset-secret") + session.add(key) + session.flush() + with ( + app.test_request_context("/"), + patch("controllers.console.datasets.datasets.ApiTokenCache.delete") as delete_cache, + ): response, status = method(api, session, "tenant-1", "api-key-id") assert status == 204 assert response == "" + delete_cache.assert_called_once() + session.flush() + assert session.get(ApiToken, "api-key-id") is None def test_delete_key_not_found(self, app: Flask): api = DatasetApiDeleteApi() method = unwrap(api.delete) - session = MagicMock() - session.scalar.return_value = None + session = self.session with app.test_request_context("/"): with pytest.raises(NotFound): method(api, session, "tenant-1", "api-key-id") -class TestDatasetEnableApiApi: +class TestDatasetEnableApiApi(_UsesSQLiteSession): @pytest.mark.parametrize(("status_value", "enabled"), [("enable", True), ("disable", False)]) def test_update_api_status(self, app: Flask, status_value: str, enabled: bool): api = DatasetEnableApiApi() method = unwrap(api.post) dataset = make_dataset(id="dataset-1") current_user = make_account() - session = MagicMock() + session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, @@ -1499,7 +1520,7 @@ class TestDatasetEnableApiApi: api = DatasetEnableApiApi() method = unwrap(api.post) dataset = make_dataset(id="dataset-1") - session = MagicMock() + session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), @@ -1582,7 +1603,7 @@ class TestDatasetRetrievalSettingApi: ] -class TestDatasetRetrievalSettingMockApi: +class TestDatasetRetrievalSettingMockApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetRetrievalSettingMockApi() method = unwrap(api.get) @@ -1597,14 +1618,14 @@ class TestDatasetRetrievalSettingMockApi: assert response["retrieval_method"] == ["semantic"] -class TestDatasetErrorDocs: +class TestDatasetErrorDocs(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetErrorDocs() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") error_doc = make_document_status(id="error-doc", indexing_status=IndexingStatus.ERROR, error="failed") current_user = make_account() - session = MagicMock() + session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, @@ -1624,7 +1645,7 @@ class TestDatasetErrorDocs: def test_get_dataset_not_found(self, app: Flask): api = DatasetErrorDocs() method = unwrap(api.get) - session = MagicMock() + session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset, @@ -1634,7 +1655,7 @@ class TestDatasetErrorDocs: get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) -class TestDatasetPermissionUserListApi: +class TestDatasetPermissionUserListApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetPermissionUserListApi() method = unwrap(api.get) @@ -1649,7 +1670,7 @@ class TestDatasetPermissionUserListApi: return_value=users, ), ): - response, status = method(api, MagicMock(), make_account(), "dataset-1") + response, status = method(api, self.session, make_account(), "dataset-1") assert status == 200 assert response["data"] == users @@ -1666,17 +1687,17 @@ class TestDatasetPermissionUserListApi: ), ): with pytest.raises(Forbidden): - method(api, MagicMock(), make_account(), "dataset-1") + method(api, self.session, make_account(), "dataset-1") -class TestDatasetAutoDisableLogApi: +class TestDatasetAutoDisableLogApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetAutoDisableLogApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") logs = {"document_ids": ["doc-1"], "count": 1} current_user = make_account() - session = MagicMock() + session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, @@ -1693,7 +1714,7 @@ class TestDatasetAutoDisableLogApi: def test_get_dataset_not_found(self, app: Flask): api = DatasetAutoDisableLogApi() method = unwrap(api.get) - session = MagicMock() + session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset, From f8f71fdc6837ca5f5cf15134201400fd5df321de Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 31 Aug 2026 05:55:38 +0000 Subject: [PATCH 13/21] test: migrate audit findings sessions and ORM models to SQLite (#40551) --- .../chat/test_base_app_runner_multimodal.py | 316 +++++++---------- .../extensions/otel/test_retrieval_tracing.py | 8 +- .../test_dataset_service_lock_not_owned.py | 49 ++- .../tasks/test_resume_agent_app_task.py | 328 ++++++++++++------ 4 files changed, 392 insertions(+), 309 deletions(-) diff --git a/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py b/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py index 6d90aa7e53b..c41c33487eb 100644 --- a/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py +++ b/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py @@ -4,12 +4,16 @@ from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest +from sqlalchemy import func, select +from sqlalchemy.orm import Session from core.app.apps.base_app_runner import AppRunner from core.app.entities.app_invoke_entities import InvokeFrom from graphon.file import FileTransferMethod, FileType from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent from models.enums import CreatorUserRole +from models.model import MessageFile +from models.tools import ToolFile class TestBaseAppRunnerMultimodal: @@ -38,18 +42,18 @@ class TestBaseAppRunnerMultimodal: return manager @pytest.fixture - def mock_tool_file(self): - """Create a mock tool file.""" - tool_file = MagicMock() - tool_file.id = str(uuid4()) - return tool_file - - @pytest.fixture - def mock_message_file(self): - """Create a mock message file.""" - message_file = MagicMock() - message_file.id = str(uuid4()) - return message_file + def tool_file(self, mock_user_id: str, mock_tenant_id: str) -> ToolFile: + """Create a real transient tool-file model returned by the external file manager.""" + return ToolFile( + user_id=mock_user_id, + tenant_id=mock_tenant_id, + conversation_id=None, + file_key="generated/image.png", + mimetype="image/png", + original_url="http://example.com/image.png", + name="image.png", + size=68, + ) def test_handle_multimodal_image_content_with_url( self, @@ -57,8 +61,8 @@ class TestBaseAppRunnerMultimodal: mock_tenant_id, mock_message_id, mock_queue_manager, - mock_tool_file, - mock_message_file, + tool_file, + sqlite_session: Session, ): """Test handling image from URL.""" # Arrange @@ -72,48 +76,33 @@ class TestBaseAppRunnerMultimodal: with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: # Setup mock tool file manager mock_mgr = MagicMock() - mock_mgr.create_file_by_url.return_value = mock_tool_file + mock_mgr.create_file_by_url.return_value = tool_file mock_mgr_class.return_value = mock_mgr - with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - # Setup mock message file - mock_msg_file_class.return_value = mock_message_file + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - file_session = MagicMock() - # Act - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) - - message_file_id = runner._handle_multimodal_image_content( - session=file_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - # Assert - mock_mgr.create_file_by_url.assert_called_once_with( - user_id=mock_user_id, - tenant_id=mock_tenant_id, - file_url=image_url, - conversation_id=None, - ) - - mock_msg_file_class.assert_called_once() - call_kwargs = mock_msg_file_class.call_args[1] - assert call_kwargs["message_id"] == mock_message_id - assert call_kwargs["type"] == FileType.IMAGE - assert call_kwargs["transfer_method"] == FileTransferMethod.TOOL_FILE - assert call_kwargs["belongs_to"] == "assistant" - assert call_kwargs["created_by"] == mock_user_id - - file_session.add.assert_called_once_with(mock_message_file) - file_session.flush.assert_called_once() - assert message_file_id == mock_message_file.id - mock_queue_manager.publish.assert_not_called() + mock_mgr.create_file_by_url.assert_called_once_with( + user_id=mock_user_id, + tenant_id=mock_tenant_id, + file_url=image_url, + conversation_id=None, + ) + message_file = sqlite_session.get(MessageFile, message_file_id) + assert message_file is not None + assert message_file.message_id == mock_message_id + assert message_file.type == FileType.IMAGE + assert message_file.transfer_method == FileTransferMethod.TOOL_FILE + assert message_file.belongs_to == "assistant" + assert message_file.created_by == mock_user_id + assert message_file.upload_file_id == tool_file.id + mock_queue_manager.publish.assert_not_called() def test_handle_multimodal_image_content_with_base64( self, @@ -121,8 +110,8 @@ class TestBaseAppRunnerMultimodal: mock_tenant_id, mock_message_id, mock_queue_manager, - mock_tool_file, - mock_message_file, + tool_file, + sqlite_session: Session, ): """Test handling image from base64 data.""" # Arrange @@ -141,41 +130,29 @@ class TestBaseAppRunnerMultimodal: with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: # Setup mock tool file manager mock_mgr = MagicMock() - mock_mgr.create_file_by_raw.return_value = mock_tool_file + mock_mgr.create_file_by_raw.return_value = tool_file mock_mgr_class.return_value = mock_mgr - with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - mock_msg_file_class.return_value = mock_message_file + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - file_session = MagicMock() - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) - - message_file_id = runner._handle_multimodal_image_content( - session=file_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - mock_mgr.create_file_by_raw.assert_called_once() - call_kwargs = mock_mgr.create_file_by_raw.call_args[1] - assert call_kwargs["user_id"] == mock_user_id - assert call_kwargs["tenant_id"] == mock_tenant_id - assert call_kwargs["conversation_id"] is None - assert "file_binary" in call_kwargs - assert call_kwargs["mimetype"] == "image/png" - assert call_kwargs["filename"].startswith("generated_image") - assert call_kwargs["filename"].endswith(".png") - - mock_msg_file_class.assert_called_once() - file_session.add.assert_called_once() - file_session.flush.assert_called_once() - assert message_file_id == mock_message_file.id - mock_queue_manager.publish.assert_not_called() + mock_mgr.create_file_by_raw.assert_called_once() + call_kwargs = mock_mgr.create_file_by_raw.call_args[1] + assert call_kwargs["user_id"] == mock_user_id + assert call_kwargs["tenant_id"] == mock_tenant_id + assert call_kwargs["conversation_id"] is None + assert "file_binary" in call_kwargs + assert call_kwargs["mimetype"] == "image/png" + assert call_kwargs["filename"].startswith("generated_image") + assert call_kwargs["filename"].endswith(".png") + assert sqlite_session.get(MessageFile, message_file_id) is not None + mock_queue_manager.publish.assert_not_called() def test_handle_multimodal_image_content_with_base64_data_uri( self, @@ -183,8 +160,8 @@ class TestBaseAppRunnerMultimodal: mock_tenant_id, mock_message_id, mock_queue_manager, - mock_tool_file, - mock_message_file, + tool_file, + sqlite_session: Session, ): """Test handling image from base64 data with URI prefix.""" # Arrange @@ -201,29 +178,22 @@ class TestBaseAppRunnerMultimodal: with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: # Setup mock tool file manager mock_mgr = MagicMock() - mock_mgr.create_file_by_raw.return_value = mock_tool_file + mock_mgr.create_file_by_raw.return_value = tool_file mock_mgr_class.return_value = mock_mgr - with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - mock_msg_file_class.return_value = mock_message_file + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - file_session = MagicMock() - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) - - runner._handle_multimodal_image_content( - session=file_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - mock_mgr.create_file_by_raw.assert_called_once() - call_kwargs = mock_mgr.create_file_by_raw.call_args[1] - assert "file_binary" in call_kwargs + mock_mgr.create_file_by_raw.assert_called_once() + call_kwargs = mock_mgr.create_file_by_raw.call_args[1] + assert "file_binary" in call_kwargs + assert sqlite_session.get(MessageFile, message_file_id) is not None def test_handle_multimodal_image_content_without_url_or_base64( self, @@ -231,6 +201,7 @@ class TestBaseAppRunnerMultimodal: mock_tenant_id, mock_message_id, mock_queue_manager, + sqlite_session: Session, ): """Test handling image content without URL or base64 data.""" # Arrange @@ -242,24 +213,19 @@ class TestBaseAppRunnerMultimodal: ) with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - file_session = MagicMock() - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) + result = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - runner._handle_multimodal_image_content( - session=file_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - mock_mgr_class.assert_not_called() - mock_msg_file_class.assert_not_called() - mock_queue_manager.publish.assert_not_called() + assert result is None + assert sqlite_session.scalar(select(func.count()).select_from(MessageFile)) == 0 + mock_mgr_class.assert_not_called() + mock_queue_manager.publish.assert_not_called() def test_handle_multimodal_image_content_with_error( self, @@ -267,6 +233,7 @@ class TestBaseAppRunnerMultimodal: mock_tenant_id, mock_message_id, mock_queue_manager, + sqlite_session: Session, ): """Test handling image content when an error occurs.""" # Arrange @@ -282,23 +249,18 @@ class TestBaseAppRunnerMultimodal: mock_mgr.create_file_by_url.side_effect = Exception("Network error") mock_mgr_class.return_value = mock_mgr - with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - file_session = MagicMock() - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) + result = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - runner._handle_multimodal_image_content( - session=file_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - mock_msg_file_class.assert_not_called() - mock_queue_manager.publish.assert_not_called() + assert result is None + assert sqlite_session.scalar(select(func.count()).select_from(MessageFile)) == 0 + mock_queue_manager.publish.assert_not_called() def test_handle_multimodal_image_content_debugger_mode( self, @@ -306,8 +268,8 @@ class TestBaseAppRunnerMultimodal: mock_tenant_id, mock_message_id, mock_queue_manager, - mock_tool_file, - mock_message_file, + tool_file, + sqlite_session: Session, ): """Test that debugger mode sets correct created_by_role.""" # Arrange @@ -321,28 +283,21 @@ class TestBaseAppRunnerMultimodal: with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: mock_mgr = MagicMock() - mock_mgr.create_file_by_url.return_value = mock_tool_file + mock_mgr.create_file_by_url.return_value = tool_file mock_mgr_class.return_value = mock_mgr - with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - mock_msg_file_class.return_value = mock_message_file + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - file_session = MagicMock() - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) - - runner._handle_multimodal_image_content( - session=file_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - call_kwargs = mock_msg_file_class.call_args[1] - assert call_kwargs["created_by_role"] == CreatorUserRole.ACCOUNT + message_file = sqlite_session.get(MessageFile, message_file_id) + assert message_file is not None + assert message_file.created_by_role == CreatorUserRole.ACCOUNT def test_handle_multimodal_image_content_service_api_mode( self, @@ -350,8 +305,8 @@ class TestBaseAppRunnerMultimodal: mock_tenant_id, mock_message_id, mock_queue_manager, - mock_tool_file, - mock_message_file, + tool_file, + sqlite_session: Session, ): """Test that service API mode sets correct created_by_role.""" # Arrange @@ -365,25 +320,18 @@ class TestBaseAppRunnerMultimodal: with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: mock_mgr = MagicMock() - mock_mgr.create_file_by_url.return_value = mock_tool_file + mock_mgr.create_file_by_url.return_value = tool_file mock_mgr_class.return_value = mock_mgr - with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - mock_msg_file_class.return_value = mock_message_file + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - file_session = MagicMock() - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) - - runner._handle_multimodal_image_content( - session=file_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - call_kwargs = mock_msg_file_class.call_args[1] - assert call_kwargs["created_by_role"] == CreatorUserRole.END_USER + message_file = sqlite_session.get(MessageFile, message_file_id) + assert message_file is not None + assert message_file.created_by_role == CreatorUserRole.END_USER diff --git a/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py b/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py index 09f0d9dc4ff..bdd418e47e4 100644 --- a/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py +++ b/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py @@ -1,10 +1,11 @@ import threading from collections.abc import Callable -from unittest.mock import MagicMock, patch +from unittest.mock import patch from uuid import uuid4 import pytest from opentelemetry.trace import StatusCode, get_current_span, get_tracer +from sqlalchemy.orm import Session from core.rag.rerank.rerank_type import RerankMode from core.rag.retrieval.dataset_retrieval import DatasetRetrieval @@ -20,6 +21,7 @@ def _otel_enabled(config_overrides: Callable[..., None]) -> None: def test_knowledge_retrieval_creates_a_child_otel_span( memory_span_exporter, tracer_provider_with_memory_exporter, + sqlite_session: Session, ) -> None: """The retrieval entry point must be visible beneath its workflow node span.""" request = KnowledgeRetrievalRequest( @@ -38,7 +40,7 @@ def test_knowledge_retrieval_creates_a_child_otel_span( patch.object(retrieval, "_get_available_datasets", return_value=[]), get_tracer(__name__).start_as_current_span("knowledge-retrieval-node") as node_span, ): - assert retrieval.knowledge_retrieval(MagicMock(), request) == [] + assert retrieval.knowledge_retrieval(sqlite_session, request) == [] retrieval_span = next( span @@ -101,7 +103,6 @@ def test_retriever_thread_exception_sets_error_span_and_is_collected( expected_error = RuntimeError("retrieval failed") with ( - patch("core.rag.retrieval.dataset_retrieval.session_factory.create_session"), patch.object(retrieval, "_retriever", side_effect=expected_error), ): retrieval._run_retriever_thread_safely( @@ -139,7 +140,6 @@ def test_retriever_thread_exception_emits_skip_event_when_requested( dataset_id = str(uuid4()) with ( - patch("core.rag.retrieval.dataset_retrieval.session_factory.create_session"), patch.object(retrieval, "_retriever", side_effect=expected_error), get_tracer(__name__).start_as_current_span("dataset-retrieval-parent") as parent_span, ): diff --git a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py index 50ca483b976..76dc9f584d2 100644 --- a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py +++ b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py @@ -1,5 +1,5 @@ import types -from unittest.mock import Mock, create_autospec +from unittest.mock import Mock import pytest from redis.exceptions import LockNotOwnedError @@ -203,19 +203,48 @@ def test_add_segment_ignores_lock_not_owned( # --------------------------------------------------------------------------- +@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, Dataset, Document, DocumentSegment)], indirect=True) def test_multi_create_segment_ignores_lock_not_owned( monkeypatch: pytest.MonkeyPatch, fake_current_user, fake_lock, + sqlite_session: Session, ): # Arrange - dataset = create_autospec(Dataset, instance=True) - dataset.id = "ds-1" - dataset.tenant_id = fake_current_user.current_tenant_id - dataset.indexing_technique = IndexTechniqueType.ECONOMY # again, skip high_quality path + dataset = Dataset( + id=DATASET_ID, + tenant_id=TENANT_ID, + name="Test Dataset", + description="", + created_by=USER_ID, + indexing_technique=IndexTechniqueType.ECONOMY, + ) + document = Document( + id=DOCUMENT_ID, + tenant_id=TENANT_ID, + dataset_id=DATASET_ID, + position=1, + data_source_type="upload_file", + data_source_info="{}", + batch="batch-1", + name="Test Document", + created_from="web", + created_by=USER_ID, + word_count=0, + doc_form=IndexStructureType.QA_INDEX, + ) + sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, document]) + sqlite_session.commit() - document = create_autospec(Document, instance=True) - document.id = "doc-1" - document.dataset_id = dataset.id - document.word_count = 0 - document.doc_form = IndexStructureType.QA_INDEX + result = SegmentService.multi_create_segment( + segments=[{"content": "question", "answer": "answer", "keywords": ["key"]}], + document=document, + dataset=dataset, + session=sqlite_session, + ) + + assert result is None + assert not sqlite_session.in_transaction() + assert sqlite_session.scalar(select(func.count(DocumentSegment.id))) == 0 + sqlite_session.refresh(document) + assert document.word_count == 0 diff --git a/api/tests/unit_tests/tasks/test_resume_agent_app_task.py b/api/tests/unit_tests/tasks/test_resume_agent_app_task.py index 684ee06edf3..c9d20849af7 100644 --- a/api/tests/unit_tests/tasks/test_resume_agent_app_task.py +++ b/api/tests/unit_tests/tasks/test_resume_agent_app_task.py @@ -1,154 +1,260 @@ -"""Unit tests for the ``resume_agent_app_execution`` celery task (ENG-635). - -Every DB access (``db.session.get``) and the generator are patched at the module -level, so the task's branch logic is exercised without a database or live stack. -""" +"""Unit tests for the ``resume_agent_app_execution`` Celery task (ENG-635).""" from __future__ import annotations -from unittest.mock import MagicMock +from collections.abc import Iterator +from datetime import UTC, datetime, timedelta +from uuid import uuid4 +import pytest from pytest_mock import MockerFixture +from sqlalchemy.orm import Session, scoped_session, sessionmaker from core.app.entities.app_invoke_entities import InvokeFrom -from models.account import Account +from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.enums import ConversationFromSource, EndUserType +from models.enums import InvokeFrom as StoredInvokeFrom from models.human_input import HumanInputForm -from models.model import App, Conversation, EndUser +from models.model import App, AppMode, Conversation, EndUser from tasks.app_generate import resume_agent_app_task as mod MODULE = "tasks.app_generate.resume_agent_app_task" -def _form(conversation_id: str = "conv-1", app_id: str = "app-1") -> MagicMock: - return MagicMock(conversation_id=conversation_id, app_id=app_id) +@pytest.fixture +def task_session(mocker: MockerFixture, sqlite_session_factory: sessionmaker[Session]) -> Iterator[Session]: + """Bind the task's Flask-SQLAlchemy session proxy to the shared SQLite database.""" + registry = scoped_session(sqlite_session_factory) + mocker.patch.object(mod.db, "session", registry) + session = registry() + yield session + registry.remove() -def _wire_db( - mocker: MockerFixture, +def _app(*, app_id: str, tenant_id: str) -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name="Agent app", + description="", + mode=AppMode.AGENT_CHAT, + icon_type=None, + icon=None, + icon_background=None, + enable_site=False, + enable_api=False, + ) + + +def _conversation( *, - form=None, - app=None, - conversation=None, - account=None, - end_user=None, -) -> MagicMock: - """Patch the module ``db`` so ``db.session.get(Model, id)`` dispatches by model.""" - table = { - HumanInputForm: form, - App: app, - Conversation: conversation, - Account: account, - EndUser: end_user, - } - db = mocker.patch(f"{MODULE}.db") - db.session.get.side_effect = lambda model, _id: table.get(model) - return db + conversation_id: str, + app_id: str, + account_id: str | None = None, + end_user_id: str | None = None, + invoke_from: StoredInvokeFrom = StoredInvokeFrom.WEB_APP, +) -> Conversation: + return Conversation( + id=conversation_id, + app_id=app_id, + mode=AppMode.AGENT_CHAT, + name="Agent conversation", + inputs={}, + invoke_from=invoke_from, + from_source=ConversationFromSource.API, + from_account_id=account_id, + from_end_user_id=end_user_id, + ) -def test_resume_happy_path_account_user_sets_tenant_and_runs(mocker: MockerFixture): - conversation = MagicMock(from_account_id="acct-1", from_end_user_id=None, invoke_from=InvokeFrom.WEB_APP) - account = MagicMock() - app = MagicMock(tenant_id="tenant-1") - db = _wire_db(mocker, form=_form(), app=app, conversation=conversation, account=account) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") - - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") - - account.set_tenant_id_with_session.assert_called_once_with("tenant-1", session=db.session.return_value) - gen.return_value.resume_after_form_submission.assert_called_once() - kwargs = gen.return_value.resume_after_form_submission.call_args.kwargs - assert kwargs["conversation_id"] == "conv-1" - assert kwargs["form_id"] == "form-1" - assert kwargs["user"] is account - assert kwargs["app_model"] is app - assert kwargs["invoke_from"] == InvokeFrom.WEB_APP - assert kwargs["session"] is db.session.return_value +def _form(*, form_id: str, conversation_id: str, app_id: str) -> HumanInputForm: + return HumanInputForm( + id=form_id, + tenant_id=str(uuid4()), + app_id=app_id, + workflow_run_id=None, + conversation_id=conversation_id, + form_kind=HumanInputFormKind.RUNTIME, + node_id="ask-human", + form_definition="{}", + rendered_content="Question", + status=HumanInputFormStatus.WAITING, + expiration_time=datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1), + ) -def test_resume_end_user_path(mocker: MockerFixture): - conversation = MagicMock(from_account_id=None, from_end_user_id="eu-1", invoke_from=InvokeFrom.WEB_APP) - end_user = MagicMock() - _wire_db(mocker, form=_form(), app=MagicMock(tenant_id="t"), conversation=conversation, end_user=end_user) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") - - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") - - assert gen.return_value.resume_after_form_submission.call_args.kwargs["user"] is end_user +def _seed_account(session: Session, *, tenant_id: str, account_id: str) -> Account: + tenant = Tenant(name="Tenant") + tenant.id = tenant_id + account = Account(name="Account", email="account@example.com") + account.id = account_id + join = TenantAccountJoin( + tenant_id=tenant_id, + account_id=account_id, + current=True, + role=TenantAccountRole.NORMAL, + ) + session.add_all([tenant, account, join]) + return account -def test_resume_preserves_debugger_invoke_from(mocker: MockerFixture): - conversation = MagicMock(from_account_id="acct-1", from_end_user_id=None, invoke_from=InvokeFrom.DEBUGGER) - account = MagicMock() - app = MagicMock(tenant_id="tenant-1") - _wire_db(mocker, form=_form(), app=app, conversation=conversation, account=account) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") +def test_resume_happy_path_account_user_sets_tenant_and_runs(mocker: MockerFixture, task_session: Session) -> None: + tenant_id, app_id, conversation_id, form_id, account_id = (str(uuid4()) for _ in range(5)) + app = _app(app_id=app_id, tenant_id=tenant_id) + account = _seed_account(task_session, tenant_id=tenant_id, account_id=account_id) + conversation = _conversation(conversation_id=conversation_id, app_id=app_id, account_id=account_id) + task_session.add_all([app, conversation, _form(form_id=form_id, conversation_id=conversation_id, app_id=app_id)]) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) - assert gen.return_value.resume_after_form_submission.call_args.kwargs["invoke_from"] == InvokeFrom.DEBUGGER + call = generator.return_value.resume_after_form_submission.call_args + assert call is not None + assert call.kwargs["conversation_id"] == conversation_id + assert call.kwargs["form_id"] == form_id + assert call.kwargs["user"] is account + assert call.kwargs["app_model"] is app + assert call.kwargs["invoke_from"] == InvokeFrom.WEB_APP + assert isinstance(call.kwargs["session"], Session) + assert account.current_tenant_id == tenant_id -def test_resume_returns_when_form_missing(mocker: MockerFixture): - _wire_db(mocker, form=None) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") +def test_resume_end_user_path(mocker: MockerFixture, task_session: Session) -> None: + tenant_id, app_id, conversation_id, form_id, end_user_id = (str(uuid4()) for _ in range(5)) + app = _app(app_id=app_id, tenant_id=tenant_id) + end_user = EndUser( + id=end_user_id, + tenant_id=tenant_id, + app_id=app_id, + type=EndUserType.BROWSER, + name="End user", + session_id="browser-session", + ) + task_session.add_all( + [ + app, + end_user, + _conversation(conversation_id=conversation_id, app_id=app_id, end_user_id=end_user_id), + _form(form_id=form_id, conversation_id=conversation_id, app_id=app_id), + ] + ) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) - gen.assert_not_called() + assert generator.return_value.resume_after_form_submission.call_args.kwargs["user"] is end_user -def test_resume_returns_on_conversation_mismatch(mocker: MockerFixture): - _wire_db(mocker, form=_form(conversation_id="other-conv")) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") +def test_resume_preserves_debugger_invoke_from(mocker: MockerFixture, task_session: Session) -> None: + tenant_id, app_id, conversation_id, form_id, account_id = (str(uuid4()) for _ in range(5)) + app = _app(app_id=app_id, tenant_id=tenant_id) + _seed_account(task_session, tenant_id=tenant_id, account_id=account_id) + task_session.add_all( + [ + app, + _conversation( + conversation_id=conversation_id, + app_id=app_id, + account_id=account_id, + invoke_from=StoredInvokeFrom.DEBUGGER, + ), + _form(form_id=form_id, conversation_id=conversation_id, app_id=app_id), + ] + ) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) - gen.assert_not_called() + assert generator.return_value.resume_after_form_submission.call_args.kwargs["invoke_from"] == InvokeFrom.DEBUGGER -def test_resume_returns_when_app_missing(mocker: MockerFixture): - _wire_db(mocker, form=_form(), app=None) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") - - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") - - gen.assert_not_called() +@pytest.mark.usefixtures("task_session") +def test_resume_returns_when_form_missing(mocker: MockerFixture) -> None: + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") + mod.resume_agent_app_execution(conversation_id=str(uuid4()), form_id=str(uuid4())) + generator.assert_not_called() -def test_resume_returns_when_conversation_missing(mocker: MockerFixture): - _wire_db(mocker, form=_form(), app=MagicMock(), conversation=None) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") - - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") - - gen.assert_not_called() +def test_resume_returns_on_conversation_mismatch(mocker: MockerFixture, task_session: Session) -> None: + app_id, form_id = str(uuid4()), str(uuid4()) + task_session.add(_form(form_id=form_id, conversation_id=str(uuid4()), app_id=app_id)) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") + mod.resume_agent_app_execution(conversation_id=str(uuid4()), form_id=form_id) + generator.assert_not_called() -def test_resume_returns_when_no_user_resolvable(mocker: MockerFixture): - conversation = MagicMock(from_account_id=None, from_end_user_id=None, invoke_from=InvokeFrom.WEB_APP) - _wire_db(mocker, form=_form(), app=MagicMock(), conversation=conversation) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") - - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") - - gen.assert_not_called() +def test_resume_returns_when_app_missing(mocker: MockerFixture, task_session: Session) -> None: + conversation_id, form_id = str(uuid4()), str(uuid4()) + task_session.add(_form(form_id=form_id, conversation_id=conversation_id, app_id=str(uuid4()))) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) + generator.assert_not_called() -def test_resume_returns_when_account_id_set_but_account_gone(mocker: MockerFixture): - conversation = MagicMock(from_account_id="acct-x", from_end_user_id=None, invoke_from=InvokeFrom.WEB_APP) - _wire_db(mocker, form=_form(), app=MagicMock(), conversation=conversation, account=None) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") - - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") - - gen.assert_not_called() +def test_resume_returns_when_conversation_missing(mocker: MockerFixture, task_session: Session) -> None: + tenant_id, app_id, conversation_id, form_id = (str(uuid4()) for _ in range(4)) + task_session.add_all( + [ + _app(app_id=app_id, tenant_id=tenant_id), + _form(form_id=form_id, conversation_id=conversation_id, app_id=app_id), + ] + ) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) + generator.assert_not_called() -def test_resume_swallows_generator_exception(mocker: MockerFixture): - conversation = MagicMock(from_account_id="acct-1", from_end_user_id=None, invoke_from=InvokeFrom.WEB_APP) - _wire_db(mocker, form=_form(), app=MagicMock(tenant_id="t"), conversation=conversation, account=MagicMock()) - gen = mocker.patch(f"{MODULE}.AgentAppGenerator") - gen.return_value.resume_after_form_submission.side_effect = RuntimeError("boom") +def test_resume_returns_when_no_user_resolvable(mocker: MockerFixture, task_session: Session) -> None: + tenant_id, app_id, conversation_id, form_id = (str(uuid4()) for _ in range(4)) + task_session.add_all( + [ + _app(app_id=app_id, tenant_id=tenant_id), + _conversation(conversation_id=conversation_id, app_id=app_id), + _form(form_id=form_id, conversation_id=conversation_id, app_id=app_id), + ] + ) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) + generator.assert_not_called() - # The task must not propagate the failure (it is logged and the session closed). - mod.resume_agent_app_execution(conversation_id="conv-1", form_id="form-1") + +def test_resume_returns_when_account_id_set_but_account_gone(mocker: MockerFixture, task_session: Session) -> None: + tenant_id, app_id, conversation_id, form_id = (str(uuid4()) for _ in range(4)) + task_session.add_all( + [ + _app(app_id=app_id, tenant_id=tenant_id), + _conversation(conversation_id=conversation_id, app_id=app_id, account_id=str(uuid4())), + _form(form_id=form_id, conversation_id=conversation_id, app_id=app_id), + ] + ) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) + generator.assert_not_called() + + +def test_resume_swallows_generator_exception(mocker: MockerFixture, task_session: Session) -> None: + tenant_id, app_id, conversation_id, form_id, account_id = (str(uuid4()) for _ in range(5)) + _seed_account(task_session, tenant_id=tenant_id, account_id=account_id) + task_session.add_all( + [ + _app(app_id=app_id, tenant_id=tenant_id), + _conversation(conversation_id=conversation_id, app_id=app_id, account_id=account_id), + _form(form_id=form_id, conversation_id=conversation_id, app_id=app_id), + ] + ) + task_session.commit() + generator = mocker.patch(f"{MODULE}.AgentAppGenerator") + generator.return_value.resume_after_form_submission.side_effect = RuntimeError("boom") + + mod.resume_agent_app_execution(conversation_id=conversation_id, form_id=form_id) + + generator.return_value.resume_after_form_submission.assert_called_once() From 1952cff091ec85ad929209d9953e2f2bee217ff2 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 31 Aug 2026 05:57:40 +0000 Subject: [PATCH 14/21] test: migrate model and agent sessions to SQLite (#40087) --- .../services/agent/test_agent_dsl_service.py | 196 +++++++------- .../agent/test_home_snapshot_service.py | 58 ++-- .../agent/test_workflow_publish_service.py | 247 +++++++++++------- .../services/agent/test_workspace_service.py | 36 +-- .../services/test_app_generate_service.py | 81 +++--- 5 files changed, 347 insertions(+), 271 deletions(-) diff --git a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py b/api/tests/unit_tests/services/agent/test_agent_dsl_service.py index 3cf2bcd6372..7c0d5ef5332 100644 --- a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py +++ b/api/tests/unit_tests/services/agent/test_agent_dsl_service.py @@ -4,10 +4,14 @@ from unittest.mock import Mock import pytest from pydantic import ValidationError +from sqlalchemy import select +from sqlalchemy.orm import Session from graphon.enums import BuiltinNodeTypes from models.agent import ( Agent, + AgentConfigDraft, + AgentConfigDraftType, AgentConfigRevision, AgentConfigRevisionOperation, AgentConfigSnapshot, @@ -191,7 +195,9 @@ def test_agent_package_rejects_null_file_id_for_available_assets(asset: dict) -> AgentPackage.model_validate(package) -def test_import_warnings_cover_runtime_setup_removed_from_package(monkeypatch: pytest.MonkeyPatch) -> None: +def test_import_warnings_cover_runtime_setup_removed_from_package( + monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: soul = AgentSoulConfig.model_validate( { "tools": { @@ -211,7 +217,7 @@ def test_import_warnings_cover_runtime_setup_removed_from_package(monkeypatch: p ) monkeypatch.setattr("services.agent.dsl_service.get_tenant_knowledge_dataset_rows", Mock(return_value={})) - _, warnings = AgentDslService(Mock())._resolve_package_soul( + _, warnings = AgentDslService(unbound_session)._resolve_package_soul( tenant_id="tenant-1", package=make_portable_agent_package(_agent(), soul), package_path="agent_packages.agent_1", @@ -231,23 +237,29 @@ def test_agent_package_rejects_unknown_schema_version() -> None: AgentPackage.model_validate(package) -def test_export_agent_app_requires_backing_agent() -> None: - session = Mock() - session.scalar.return_value = None - +def test_export_agent_app_requires_backing_agent(sqlite_session: Session) -> None: with pytest.raises(ValueError, match="no active backing Agent"): - AgentDslService(session).export_agent_app(app=SimpleNamespace(tenant_id="tenant-1", id="app-1")) + AgentDslService(sqlite_session).export_agent_app(app=SimpleNamespace(tenant_id="tenant-1", id="app-1")) @pytest.mark.parametrize("use_draft", [True, False]) -def test_export_agent_app_uses_draft_or_active_snapshot(use_draft: bool) -> None: +def test_export_agent_app_uses_draft_or_active_snapshot(use_draft: bool, sqlite_session: Session) -> None: agent = _agent() + agent.app_id = "app-1" agent.active_config_snapshot_id = "snapshot-1" - draft = SimpleNamespace(config_snapshot_dict=AgentSoulConfig(config_note="draft").model_dump(mode="json")) - session = Mock() - session.scalar.side_effect = [agent, draft if use_draft else None] - session.execute.return_value = [] - service = AgentDslService(session) + sqlite_session.add(agent) + if use_draft: + sqlite_session.add( + AgentConfigDraft( + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + config_snapshot=AgentSoulConfig(config_note="draft"), + ) + ) + sqlite_session.flush() + service = AgentDslService(sqlite_session) require_snapshot = Mock(return_value=_snapshot(soul=AgentSoulConfig(config_note="snapshot"))) service._require_snapshot = require_snapshot @@ -258,22 +270,26 @@ def test_export_agent_app_uses_draft_or_active_snapshot(use_draft: bool) -> None assert require_snapshot.call_count == (0 if use_draft else 1) -def test_export_workflow_packages_deduplicates_shared_agent() -> None: +def test_export_workflow_packages_deduplicates_shared_agent(sqlite_session: Session) -> None: graph = {"nodes": [_agent_node("node-1"), _agent_node("node-2")], "edges": []} bindings = [ - SimpleNamespace( + WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version="draft", node_id=node_id, agent_id="agent-1", current_snapshot_id="snapshot-1", binding_type=WorkflowAgentBindingType.ROSTER_AGENT, - node_job_config_dict={"workflow_prompt": node_id}, + node_job_config={"workflow_prompt": node_id}, + created_by="account-1", ) for node_id in ("node-1", "node-2") ] - session = Mock() - session.scalars.return_value.all.return_value = bindings - session.execute.return_value = [] - service = AgentDslService(session) + sqlite_session.add_all(bindings) + sqlite_session.flush() + service = AgentDslService(sqlite_session) service._require_agent = Mock(return_value=_agent()) service._require_snapshot = Mock(return_value=_snapshot()) @@ -292,12 +308,9 @@ def test_export_workflow_packages_deduplicates_shared_agent() -> None: assert service._require_agent.call_count == 2 -def test_export_workflow_packages_rejects_incomplete_binding() -> None: - session = Mock() - session.scalars.return_value.all.return_value = [] - +def test_export_workflow_packages_rejects_incomplete_binding(sqlite_session: Session) -> None: with pytest.raises(ValueError, match="no complete persisted binding"): - AgentDslService(session).export_workflow_packages( + AgentDslService(sqlite_session).export_workflow_packages( workflow=SimpleNamespace(tenant_id="tenant-1", id="workflow-1", version="draft"), graph={"nodes": [_agent_node("node-1")], "edges": []}, ) @@ -328,9 +341,10 @@ def test_graph_without_package_bindings_removes_portable_fields() -> None: assert AGENT_NODE_JOB_DSL_KEY in graph["nodes"][0]["data"] -def test_import_agent_app_package_creates_config_and_unpublished_draft(monkeypatch: pytest.MonkeyPatch) -> None: - session = Mock() - service = AgentDslService(session) +def test_import_agent_app_package_creates_config_and_unpublished_draft( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + service = AgentDslService(sqlite_session) soul = AgentSoulConfig(config_note="portable") warning = DslImportWarning(code="setup", path="agent.soul", message="setup required") service._resolve_package_soul = Mock(return_value=(soul, [warning])) @@ -361,11 +375,10 @@ def test_import_agent_app_package_creates_config_and_unpublished_draft(monkeypat assert agent.active_config_is_published is False assert app.name == "Portable Agent" assert app.description == "description" - assert session.add.call_count == 2 - assert session.flush.call_count == 2 + assert sqlite_session.scalar(select(AgentConfigDraft).where(AgentConfigDraft.agent_id == agent.id)) is not None -def test_import_workflow_packages_materializes_every_package_binding_as_inline() -> None: +def test_import_workflow_packages_materializes_every_package_binding_as_inline(sqlite_session: Session) -> None: package = make_portable_agent_package(_agent(), AgentSoulConfig()) graph = { "nodes": [ @@ -388,14 +401,22 @@ def test_import_workflow_packages_materializes_every_package_binding_as_inline() } for node in graph["nodes"][:3]: node["data"][AGENT_NODE_JOB_DSL_KEY] = {"workflow_prompt": node["id"]} - old_binding = SimpleNamespace( + old_binding = WorkflowAgentNodeBinding( id="old-binding", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version="draft", + node_id="old-node", binding_type=WorkflowAgentBindingType.INLINE_AGENT, agent_id="old-inline-agent", + current_snapshot_id="old-snapshot", + node_job_config={}, + created_by="account-1", ) - session = Mock() - session.scalars.return_value.all.return_value = [old_binding] - service = AgentDslService(session) + sqlite_session.add(old_binding) + sqlite_session.flush() + service = AgentDslService(sqlite_session) imported_results = [ SimpleNamespace( agent=SimpleNamespace(id=f"inline-agent-{index}"), @@ -420,7 +441,7 @@ def test_import_workflow_packages_materializes_every_package_binding_as_inline() account=SimpleNamespace(id="account-1"), ) - session.delete.assert_called_once_with(old_binding) + assert sqlite_session.get(WorkflowAgentNodeBinding, "old-binding") is None assert retirement_candidates == {"old-inline-agent"} assert service._create_imported_inline_agent.call_count == 3 assert [call.kwargs["node_id"] for call in service._create_imported_inline_agent.call_args_list] == [ @@ -438,8 +459,10 @@ def test_import_workflow_packages_materializes_every_package_binding_as_inline() assert all(binding["binding_type"] == WorkflowAgentBindingType.INLINE_AGENT.value for binding in bindings) assert AGENT_NODE_JOB_DSL_KEY not in result["nodes"][0]["data"] assert json.loads(workflow.graph) == result - added_bindings = [item.args[0] for item in session.add.call_args_list] - assert all(isinstance(binding, WorkflowAgentNodeBinding) for binding in added_bindings) + added_bindings = sqlite_session.scalars( + select(WorkflowAgentNodeBinding).where(WorkflowAgentNodeBinding.workflow_id == "workflow-1") + ).all() + assert len(added_bindings) == 3 assert all(binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT for binding in added_bindings) @@ -453,13 +476,13 @@ def test_import_workflow_packages_materializes_every_package_binding_as_inline() ({"binding_type": "invalid", AGENT_PACKAGE_REF_KEY: "agent_1"}, "invalid binding type"), ], ) -def test_import_workflow_packages_rejects_invalid_package_binding(binding: dict, error: str) -> None: - session = Mock() - session.scalars.return_value.all.return_value = [] +def test_import_workflow_packages_rejects_invalid_package_binding( + binding: dict, error: str, sqlite_session: Session +) -> None: package = make_portable_agent_package(_agent(), AgentSoulConfig()) with pytest.raises(ValueError, match=error): - AgentDslService(session).import_workflow_packages( + AgentDslService(sqlite_session).import_workflow_packages( workflow=SimpleNamespace(tenant_id="tenant-1", app_id="app-1", id="workflow-1", version="draft"), portable_graph={"nodes": [_agent_node("node-1", binding)], "edges": []}, raw_packages={"agent_1": package.model_dump(mode="json")}, @@ -467,9 +490,8 @@ def test_import_workflow_packages_rejects_invalid_package_binding(binding: dict, ) -def test_clone_inline_binding_copies_soul() -> None: - session = Mock() - service = AgentDslService(session) +def test_clone_inline_binding_copies_soul(unbound_session: Session) -> None: + service = AgentDslService(unbound_session) target_agent = SimpleNamespace(id="target-agent") target_snapshot = SimpleNamespace(id="target-snapshot") service._create_workflow_only_agent = Mock(return_value=(target_agent, target_snapshot)) @@ -504,7 +526,9 @@ def test_clone_inline_binding_copies_soul() -> None: assert create_kwargs["source"] == AgentSource.WORKFLOW -def test_extract_package_dependencies_covers_model_tools_and_knowledge(monkeypatch: pytest.MonkeyPatch) -> None: +def test_extract_package_dependencies_covers_model_tools_and_knowledge( + monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: model_dependency = Mock(side_effect=lambda provider: f"model:{provider}") tool_dependency = Mock(side_effect=lambda provider: f"tool:{provider}") monkeypatch.setattr( @@ -551,7 +575,7 @@ def test_extract_package_dependencies_covers_model_tools_and_knowledge(monkeypat } ) - dependencies = AgentDslService(Mock()).extract_package_dependencies( + dependencies = AgentDslService(unbound_session).extract_package_dependencies( {"agent_1": make_portable_agent_package(_agent(), soul)} ) @@ -564,8 +588,8 @@ def test_extract_package_dependencies_covers_model_tools_and_knowledge(monkeypat ] -def test_create_imported_inline_agent_uses_import_provenance() -> None: - service = AgentDslService(Mock()) +def test_create_imported_inline_agent_uses_import_provenance(unbound_session: Session) -> None: + service = AgentDslService(unbound_session) soul = AgentSoulConfig(config_note="inline") warning = DslImportWarning(code="setup", path="agent", message="setup") service._resolve_package_soul = Mock(return_value=(soul, [warning])) @@ -587,9 +611,10 @@ def test_create_imported_inline_agent_uses_import_provenance() -> None: ) -def test_create_workflow_only_agent_sets_backing_app_and_snapshot(monkeypatch: pytest.MonkeyPatch) -> None: - session = Mock() - service = AgentDslService(session) +def test_create_workflow_only_agent_sets_backing_app_and_snapshot( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + service = AgentDslService(sqlite_session) roster_service = Mock() roster_service.create_hidden_backing_app_for_workflow_agent.return_value = SimpleNamespace(id="backing-app") monkeypatch.setattr("services.agent.dsl_service.AgentRosterService", Mock(return_value=roster_service)) @@ -613,11 +638,12 @@ def test_create_workflow_only_agent_sets_backing_app_and_snapshot(monkeypatch: p assert agent.active_config_snapshot_id == "snapshot-1" assert agent.active_config_has_model is True assert agent.active_config_is_published is True - session.add.assert_called_once_with(agent) - assert session.flush.call_count == 2 + assert sqlite_session.get(Agent, agent.id) is agent -def test_resolve_package_soul_preserves_existing_and_marks_missing_knowledge(monkeypatch: pytest.MonkeyPatch) -> None: +def test_resolve_package_soul_preserves_existing_and_marks_missing_knowledge( + monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: soul = AgentSoulConfig.model_validate( { "config_skills": [{"name": "skill", "file_kind": "tool_file", "file_id": "skill-file"}], @@ -638,18 +664,17 @@ def test_resolve_package_soul_preserves_existing_and_marks_missing_knowledge(mon }, } ) - session = Mock() get_dataset_rows = Mock(return_value={"existing": SimpleNamespace(id="existing")}) monkeypatch.setattr("services.agent.dsl_service.get_tenant_knowledge_dataset_rows", get_dataset_rows) - resolved, warnings = AgentDslService(session)._resolve_package_soul( + resolved, warnings = AgentDslService(unbound_session)._resolve_package_soul( tenant_id="tenant-1", package=make_portable_agent_package(_agent(), soul), package_path="agent_packages.agent_1", ) get_dataset_rows.assert_called_once_with( - session=session, + session=unbound_session, tenant_id="tenant-1", dataset_ids=["existing", "missing"], ) @@ -683,14 +708,18 @@ def test_resolve_package_soul_preserves_existing_and_marks_missing_knowledge(mon } -def test_create_snapshot_increments_version_and_records_revision() -> None: - session = Mock() - session.scalar.return_value = 2 - service = AgentDslService(session) +def test_create_snapshot_increments_version_and_records_revision(sqlite_session: Session) -> None: + agent = _agent() + first = _snapshot(snapshot_id="snapshot-1") + second = _snapshot(snapshot_id="snapshot-2") + second.version = 2 + sqlite_session.add_all([agent, first, second]) + sqlite_session.flush() + service = AgentDslService(sqlite_session) snapshot = service._create_snapshot( tenant_id="tenant-1", - agent=_agent(), + agent=agent, account_id="account-1", soul=AgentSoulConfig(config_note="version 3"), operation=AgentConfigRevisionOperation.IMPORT_PACKAGE, @@ -698,28 +727,32 @@ def test_create_snapshot_increments_version_and_records_revision() -> None: assert snapshot.version == 3 assert snapshot.home_snapshot_id is None - assert isinstance(session.add.call_args_list[0].args[0], AgentConfigSnapshot) - revision = session.add.call_args_list[1].args[0] - assert isinstance(revision, AgentConfigRevision) + revision = sqlite_session.scalar( + select(AgentConfigRevision).where(AgentConfigRevision.current_snapshot_id == snapshot.id) + ) + assert revision is not None assert revision.operation == AgentConfigRevisionOperation.IMPORT_PACKAGE - assert session.flush.call_count == 2 -def test_unique_roster_name_uses_first_available_suffix() -> None: - session = Mock() - session.scalars.return_value.all.return_value = ["Agent", "Agent import"] +def test_unique_roster_name_uses_first_available_suffix(sqlite_session: Session) -> None: + for index, name in enumerate(("Agent", "Agent import"), start=1): + agent = _agent() + agent.id = f"agent-{index}" + agent.name = name + sqlite_session.add(agent) + sqlite_session.flush() - result = AgentDslService(session)._unique_roster_name(tenant_id="tenant-1", requested="Agent") + result = AgentDslService(sqlite_session)._unique_roster_name(tenant_id="tenant-1", requested="Agent") assert result == "Agent import 2" -def test_require_helpers_and_graph_detection() -> None: - session = Mock() - service = AgentDslService(session) +def test_require_helpers_and_graph_detection(sqlite_session: Session) -> None: + service = AgentDslService(sqlite_session) agent = _agent() snapshot = _snapshot() - session.scalar.side_effect = [agent, None, snapshot, None] + sqlite_session.add_all([agent, snapshot]) + sqlite_session.flush() assert service._require_agent(tenant_id="tenant-1", agent_id="agent-1") is agent with pytest.raises(ValueError, match="source Agent"): @@ -733,17 +766,4 @@ def test_require_helpers_and_graph_detection() -> None: assert AgentDslService._agent_icon_type(AgentIconType.EMOJI.value) == AgentIconType.EMOJI assert AgentDslService._agent_icon_type(None) is None assert is_agent_v2_graph({"nodes": [_agent_node("agent")]}) is True - assert is_agent_v2_graph({"nodes": [{"id": "legacy-agent", "data": {"type": "agent", "version": "2"}}]}) is False assert is_agent_v2_graph({"nodes": ["invalid", {"data": {"type": "start"}}]}) is False - - -def test_export_workflow_packages_ignores_historical_agent_version_two() -> None: - session = Mock() - service = AgentDslService(session) - graph = {"nodes": [{"id": "legacy-agent", "data": {"type": "agent", "version": "2"}}]} - - portable_graph, packages = service.export_workflow_packages(workflow=Mock(), graph=graph) - - assert portable_graph == graph - assert packages == {} - session.scalars.assert_not_called() 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 6cb64587818..3b2a638f692 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 @@ -12,6 +12,9 @@ from models.agent import ( AgentConfigDraftType, AgentConfigSnapshot, AgentHomeSnapshot, + AgentScope, + AgentSource, + AgentStatus, AgentWorkingResourceStatus, ) from models.agent_config_entities import AgentSoulConfig @@ -53,16 +56,31 @@ def test_home_snapshot_client_outlasts_the_gateway_snapshot_budget(monkeypatch: assert client._timeout == 45.0 -def test_validate_home_snapshot_binding_accepts_default_home_without_ledger_lookup() -> None: - session = MagicMock() +def _persist_agent(session: Session, *, app_id: str, backing_app_id: str | None) -> Agent: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Snapshot Agent", + description="", + role="", + scope=AgentScope.ROSTER if backing_app_id is None else AgentScope.WORKFLOW_ONLY, + source=AgentSource.AGENT_APP if backing_app_id is None else AgentSource.WORKFLOW, + status=AgentStatus.ACTIVE, + app_id=app_id, + backing_app_id=backing_app_id, + ) + session.add(agent) + session.commit() + return agent + + +def test_validate_home_snapshot_binding_accepts_default_home_without_ledger_lookup(unbound_session: Session) -> None: validate_home_snapshot_binding( - session=session, + session=unbound_session, agent=Agent(id="agent-1"), home_snapshot_id=None, ) - session.scalar.assert_not_called() - @pytest.mark.parametrize( ("app_id", "backing_app_id", "expected_runtime_app_id"), @@ -73,12 +91,12 @@ def test_validate_home_snapshot_binding_accepts_default_home_without_ledger_look ) def test_build_apply_checkpoints_exact_active_binding( monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, app_id: str, backing_app_id: str | None, expected_runtime_app_id: str, ) -> None: - session = MagicMock() - session.scalar.return_value = SimpleNamespace(app_id=app_id, backing_app_id=backing_app_id) + _persist_agent(sqlite_session, app_id=app_id, backing_app_id=backing_app_id) binding = SimpleNamespace( backend_binding_ref="binding-ref-1", agent_id="agent-1", @@ -94,7 +112,7 @@ def test_build_apply_checkpoints_exact_active_binding( monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) snapshot = AgentHomeSnapshotService.create_for_build_apply( - session=session, + session=sqlite_session, build_draft=_build_draft(), ) @@ -106,9 +124,8 @@ def test_build_apply_checkpoints_exact_active_binding( assert validate_generation.call_args.kwargs["base_home_snapshot_id"] == "home-old" -def test_build_apply_forwards_default_home_generation(monkeypatch: pytest.MonkeyPatch) -> None: - session = MagicMock() - session.scalar.return_value = SimpleNamespace(app_id="app-1", backing_app_id=None) +def test_build_apply_forwards_default_home_generation(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + _persist_agent(sqlite_session, app_id="app-1", backing_app_id=None) binding = SimpleNamespace( backend_binding_ref="binding-ref-1", agent_id="agent-1", @@ -123,7 +140,7 @@ def test_build_apply_forwards_default_home_generation(monkeypatch: pytest.Monkey monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) snapshot = AgentHomeSnapshotService.create_for_build_apply( - session=session, + session=sqlite_session, build_draft=_build_draft(home_snapshot_id=None), ) @@ -131,26 +148,26 @@ def test_build_apply_forwards_default_home_generation(monkeypatch: pytest.Monkey assert validate_generation.call_args.kwargs["base_home_snapshot_id"] is None -def test_build_apply_fails_fast_without_source_binding() -> None: - session = MagicMock() +def test_build_apply_fails_fast_without_source_binding(unbound_session: Session) -> None: build_draft = _build_draft() build_draft.agent_workspace_binding_id = None with pytest.raises(AgentBuildSandboxNotFoundError): AgentHomeSnapshotService.create_for_build_apply( - session=session, + session=unbound_session, build_draft=build_draft, ) -def test_home_snapshot_collection_database_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: - context = MagicMock() - session = context.__enter__.return_value +def test_home_snapshot_collection_database_failure_propagates( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: error = RuntimeError("database unavailable") - session.scalar.side_effect = error + scalar = MagicMock(side_effect=error) + monkeypatch.setattr(sqlite_session, "scalar", scalar) monkeypatch.setattr( "services.agent.home_snapshot_service.session_factory.create_session", - lambda: context, + lambda: nullcontext(sqlite_session), ) with pytest.raises(RuntimeError) as exc_info: @@ -159,6 +176,7 @@ def test_home_snapshot_collection_database_failure_propagates(monkeypatch: pytes home_snapshot_id="home-1", ) + scalar.assert_called_once() assert exc_info.value is error 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 74031c57844..d50f3589b68 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 @@ -7,15 +7,19 @@ from sqlalchemy.orm import Session from models.agent import ( Agent, + AgentConfigSnapshot, AgentScope, + AgentSource, + AgentStatus, WorkflowAgentBindingType, WorkflowAgentNodeBinding, ) +from models.agent_config_entities import AgentSoulConfig from models.enums import AppStatus from models.model import App, AppMode from models.workflow import Workflow, WorkflowType from services.agent.dsl_service import AgentDslService -from services.agent.workflow_publish_service import WorkflowAgentPublishService, _InlineAgentOwnershipError +from services.agent.workflow_publish_service import WorkflowAgentPublishService def _workflow(*, workflow_id: str = "workflow-1", version: str = Workflow.VERSION_DRAFT) -> Workflow: @@ -33,39 +37,66 @@ def _workflow(*, workflow_id: str = "workflow-1", version: str = Workflow.VERSIO ) -def test_inline_binding_from_another_node_is_cloned(monkeypatch: pytest.MonkeyPatch) -> None: - session = Mock() - draft_workflow = _workflow() - monkeypatch.setattr( - WorkflowAgentPublishService, - "_resolve_inline_agent_graph_binding", - Mock(side_effect=_InlineAgentOwnershipError("source belongs to another node")), +def _inline_agent( + *, + agent_id: str, + workflow_id: str, + node_id: str, + tenant_id: str = "tenant-1", +) -> Agent: + return Agent( + id=agent_id, + tenant_id=tenant_id, + name=f"Inline {agent_id}", + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + status=AgentStatus.ACTIVE, + app_id="app-1", + workflow_id=workflow_id, + workflow_node_id=node_id, ) - monkeypatch.setattr( - WorkflowAgentPublishService, - "_resolve_existing_inline_binding_agent", - Mock(return_value=None), + + +def _snapshot(*, snapshot_id: str, agent_id: str, version: int = 1) -> AgentConfigSnapshot: + return AgentConfigSnapshot( + id=snapshot_id, + tenant_id="tenant-1", + agent_id=agent_id, + version=version, + config_snapshot=AgentSoulConfig(), ) - clone = Mock(return_value=(SimpleNamespace(id="target-agent"), "target-snapshot")) - monkeypatch.setattr(WorkflowAgentPublishService, "_clone_inline_graph_binding_for_node", clone) + + +def test_inline_binding_from_another_node_is_cloned(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + source_agent = _inline_agent(agent_id="source-agent", workflow_id="workflow-1", node_id="source-node") + source_snapshot = _snapshot(snapshot_id="source-snapshot", agent_id=source_agent.id) + sqlite_session.add_all([source_agent, source_snapshot]) + sqlite_session.commit() + target_agent = _inline_agent(agent_id="target-agent", workflow_id="workflow-1", node_id="pasted-node") + target_snapshot = _snapshot(snapshot_id="target-snapshot", agent_id=target_agent.id) + clone = Mock(return_value=(target_agent, target_snapshot)) + monkeypatch.setattr(AgentDslService, "clone_inline_binding_for_node", clone) WorkflowAgentPublishService._sync_agent_binding_for_node( - session=session, - draft_workflow=draft_workflow, + session=sqlite_session, + draft_workflow=_workflow(), node_id="pasted-node", node_data={"agent_task": "Summarize the input"}, node_binding={ "binding_type": WorkflowAgentBindingType.INLINE_AGENT.value, - "agent_id": "source-agent", - "current_snapshot_id": "source-snapshot", + "agent_id": source_agent.id, + "current_snapshot_id": source_snapshot.id, }, existing_binding=None, account_id="account-1", ) + sqlite_session.flush() clone.assert_called_once() - binding = session.add.call_args.args[0] - assert isinstance(binding, WorkflowAgentNodeBinding) + binding = sqlite_session.scalar( + select(WorkflowAgentNodeBinding).where(WorkflowAgentNodeBinding.node_id == "pasted-node") + ) + assert binding is not None assert binding.agent_id == "target-agent" assert binding.current_snapshot_id == "target-snapshot" assert binding.node_job_config.workflow_prompt == "Summarize the input" @@ -103,8 +134,9 @@ def test_draft_sync_resolves_roster_agents() -> None: 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: +def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent(sqlite_session: Session) -> None: existing_inline = WorkflowAgentNodeBinding( + id="existing-inline", tenant_id="tenant-1", app_id="app-1", workflow_id="draft-workflow", @@ -117,6 +149,7 @@ def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> N created_by="account-1", ) existing_roster = WorkflowAgentNodeBinding( + id="existing-roster", tenant_id="tenant-1", app_id="app-1", workflow_id="draft-workflow", @@ -129,6 +162,7 @@ def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> N created_by="account-1", ) source = WorkflowAgentNodeBinding( + id="source-roster", tenant_id="tenant-1", app_id="app-1", workflow_id="published-workflow", @@ -140,30 +174,34 @@ def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> N node_job_config={"workflow_prompt": "Use the roster agent"}, created_by="account-1", ) - session = Mock() - session.scalars.side_effect = [ - SimpleNamespace(all=lambda: [existing_inline, existing_roster]), - SimpleNamespace(all=lambda: [source]), - ] - session.scalar.return_value = SimpleNamespace( + roster_agent = Agent( id="roster-agent", + tenant_id="tenant-1", + name="Roster Agent", scope=AgentScope.ROSTER, + source=AgentSource.ROSTER, + status=AgentStatus.ACTIVE, + app_id="roster-app", active_config_snapshot_id="published-snapshot", ) + sqlite_session.add_all([existing_inline, existing_roster, source, roster_agent]) + sqlite_session.commit() retirement_candidates = WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( - session=session, + session=sqlite_session, source_workflow=_workflow(workflow_id="published-workflow", version="2026-07-13 00:00:00"), draft_workflow=_workflow(workflow_id="draft-workflow"), account_id="account-2", ) - assert {item.args[0].agent_id for item in session.delete.call_args_list} == { - "old-inline-agent", - "old-roster-agent", - } - restored = session.add.call_args.args[0] - assert isinstance(restored, WorkflowAgentNodeBinding) - assert restored.workflow_id == "draft-workflow" + assert sqlite_session.get(WorkflowAgentNodeBinding, existing_inline.id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, existing_roster.id) is None + restored = sqlite_session.scalar( + select(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.workflow_id == "draft-workflow", + WorkflowAgentNodeBinding.node_id == "agent-node", + ) + ) + assert restored is not None assert restored.workflow_version == Workflow.VERSION_DRAFT assert restored.agent_id == "roster-agent" assert restored.current_snapshot_id == "published-snapshot" @@ -284,6 +322,7 @@ def test_publish_binding_copy_keeps_previous_published_owner( draft_workflow=draft_workflow, published_workflow=published_workflow, ) + sqlite_session.flush() assert result is True assert sqlite_session.get(WorkflowAgentNodeBinding, previous_inline_binding.id) is previous_inline_binding @@ -299,55 +338,50 @@ def test_publish_binding_copy_keeps_previous_published_owner( assert copied.current_snapshot_id == "draft-inline-snapshot" -def test_inline_binding_reuses_existing_node_owned_agent(monkeypatch: pytest.MonkeyPatch) -> None: - session = Mock() - draft_workflow = _workflow() +def test_inline_binding_reuses_existing_node_owned_agent(sqlite_session: Session) -> None: + existing_agent = _inline_agent(agent_id="existing-agent", workflow_id="workflow-1", node_id="pasted-node") + existing_snapshot = _snapshot(snapshot_id="existing-snapshot", agent_id=existing_agent.id) existing_binding = WorkflowAgentNodeBinding( + id="existing-binding", tenant_id="tenant-1", app_id="app-1", workflow_id="workflow-1", workflow_version=Workflow.VERSION_DRAFT, node_id="pasted-node", binding_type=WorkflowAgentBindingType.INLINE_AGENT, - agent_id="existing-agent", - current_snapshot_id="existing-snapshot", + agent_id=existing_agent.id, + current_snapshot_id=existing_snapshot.id, node_job_config={}, created_by="account-1", ) - existing_agent = SimpleNamespace(id="existing-agent") - monkeypatch.setattr( - WorkflowAgentPublishService, - "_resolve_inline_agent_graph_binding", - Mock(side_effect=_InlineAgentOwnershipError("source belongs to another node")), - ) - monkeypatch.setattr( - WorkflowAgentPublishService, - "_resolve_existing_inline_binding_agent", - Mock(return_value=existing_agent), - ) - clone = Mock() - monkeypatch.setattr(WorkflowAgentPublishService, "_clone_inline_graph_binding_for_node", clone) + sqlite_session.add_all([existing_agent, existing_snapshot, existing_binding]) + sqlite_session.commit() WorkflowAgentPublishService._sync_agent_binding_for_node( - session=session, - draft_workflow=draft_workflow, + session=sqlite_session, + draft_workflow=_workflow(), node_id="pasted-node", node_data={"agent_task": "Summarize"}, node_binding={ "binding_type": WorkflowAgentBindingType.INLINE_AGENT.value, - "agent_id": "source-agent", - "current_snapshot_id": "source-snapshot", + "agent_id": "unavailable-source-agent", + "current_snapshot_id": "unavailable-source-snapshot", }, existing_binding=existing_binding, account_id="account-1", ) + sqlite_session.flush() - assert existing_binding.agent_id == "existing-agent" - assert existing_binding.current_snapshot_id == "existing-snapshot" - clone.assert_not_called() + stored = sqlite_session.get(WorkflowAgentNodeBinding, existing_binding.id) + assert stored is not None + assert stored.agent_id == "existing-agent" + assert stored.current_snapshot_id == "existing-snapshot" + assert stored.node_job_config.workflow_prompt == "Summarize" -def test_resolve_existing_inline_binding_agent_returns_valid_agent_or_none(monkeypatch: pytest.MonkeyPatch) -> None: +def test_resolve_existing_inline_binding_agent_returns_valid_agent_or_none( + monkeypatch: pytest.MonkeyPatch, unbound_session: Session +) -> None: binding = WorkflowAgentNodeBinding( tenant_id="tenant-1", app_id="app-1", @@ -360,13 +394,13 @@ def test_resolve_existing_inline_binding_agent_returns_valid_agent_or_none(monke node_job_config={}, created_by="account-1", ) - resolved = SimpleNamespace(id="agent-1") + resolved = _inline_agent(agent_id="agent-1", workflow_id="workflow-1", node_id="node-1") resolver = Mock(return_value=resolved) monkeypatch.setattr(WorkflowAgentPublishService, "_resolve_inline_agent_graph_binding", resolver) assert ( WorkflowAgentPublishService._resolve_existing_inline_binding_agent( - session=Mock(), + session=unbound_session, draft_workflow=_workflow(), node_id="node-1", existing_binding=binding, @@ -377,7 +411,7 @@ def test_resolve_existing_inline_binding_agent_returns_valid_agent_or_none(monke resolver.side_effect = ValueError("stale") assert ( WorkflowAgentPublishService._resolve_existing_inline_binding_agent( - session=Mock(), + session=unbound_session, draft_workflow=_workflow(), node_id="node-1", existing_binding=binding, @@ -386,30 +420,42 @@ 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 +def test_resolve_roster_binding_rejects_unpublished_agent(sqlite_session: Session) -> None: + sqlite_session.add( + Agent( + id="decoy-agent", + tenant_id="tenant-1", + name="Decoy", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="decoy-app", + ) + ) + sqlite_session.commit() with pytest.raises(ValueError, match="unavailable or unpublished roster agent"): WorkflowAgentPublishService._resolve_roster_agent_graph_binding( - session=session, + session=sqlite_session, draft_workflow=_workflow(), node_id="agent-node", agent_id="agent-1", ) -def test_clone_inline_graph_binding_for_node_clones_source(monkeypatch: pytest.MonkeyPatch) -> None: - session = Mock() - source_agent = SimpleNamespace(id="source-agent") - source_snapshot = SimpleNamespace(id="source-snapshot") - session.scalar.side_effect = [source_agent, source_snapshot] - target_agent = SimpleNamespace(id="target-agent") - target_snapshot = SimpleNamespace(id="target-snapshot") +def test_clone_inline_graph_binding_for_node_clones_source( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + source_agent = _inline_agent(agent_id="source-agent", workflow_id="source-workflow", node_id="source-node") + source_snapshot = _snapshot(snapshot_id="source-snapshot", agent_id=source_agent.id) + sqlite_session.add_all([source_agent, source_snapshot]) + sqlite_session.commit() + target_agent = _inline_agent(agent_id="target-agent", workflow_id="workflow-1", node_id="target-node") + target_snapshot = _snapshot(snapshot_id="target-snapshot", agent_id=target_agent.id) clone = Mock(return_value=(target_agent, target_snapshot)) monkeypatch.setattr(AgentDslService, "clone_inline_binding_for_node", clone) result = WorkflowAgentPublishService._clone_inline_graph_binding_for_node( - session=session, + session=sqlite_session, draft_workflow=_workflow(), node_id="target-node", source_agent_id="source-agent", @@ -427,14 +473,17 @@ def test_clone_inline_graph_binding_for_node_clones_source(monkeypatch: pytest.M ) -@pytest.mark.parametrize("scalar_results", [[None], [SimpleNamespace(id="source-agent"), None]]) -def test_clone_inline_graph_binding_for_node_rejects_missing_source(scalar_results: list[object | None]) -> None: - session = Mock() - session.scalar.side_effect = scalar_results +@pytest.mark.parametrize("persist_source_agent", [False, True]) +def test_clone_inline_graph_binding_for_node_rejects_missing_source( + sqlite_session: Session, persist_source_agent: bool +) -> None: + if persist_source_agent: + sqlite_session.add(_inline_agent(agent_id="source-agent", workflow_id="source-workflow", node_id="source-node")) + sqlite_session.commit() with pytest.raises(ValueError, match="unavailable inline agent|missing inline agent config snapshot"): WorkflowAgentPublishService._clone_inline_graph_binding_for_node( - session=session, + session=sqlite_session, draft_workflow=_workflow(), node_id="target-node", source_agent_id="source-agent", @@ -443,37 +492,45 @@ def test_clone_inline_graph_binding_for_node_rejects_missing_source(scalar_resul ) -def test_restore_clones_inline_binding_owned_by_published_workflow(monkeypatch: pytest.MonkeyPatch) -> None: +def test_restore_clones_inline_binding_owned_by_published_workflow( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + source_agent = _inline_agent(agent_id="published-agent", workflow_id="published-workflow", node_id="agent-node") + source_snapshot = _snapshot(snapshot_id="published-snapshot", agent_id=source_agent.id) source = WorkflowAgentNodeBinding( + id="published-binding", tenant_id="tenant-1", app_id="app-1", workflow_id="published-workflow", workflow_version="published", node_id="agent-node", binding_type=WorkflowAgentBindingType.INLINE_AGENT, - agent_id="published-agent", - current_snapshot_id="published-snapshot", + agent_id=source_agent.id, + current_snapshot_id=source_snapshot.id, node_job_config={"workflow_prompt": "work"}, created_by="account-1", ) - session = Mock() - session.scalars.side_effect = [SimpleNamespace(all=lambda: []), SimpleNamespace(all=lambda: [source])] - monkeypatch.setattr( - WorkflowAgentPublishService, - "_resolve_inline_agent_graph_binding", - Mock(side_effect=ValueError("owned by published workflow")), - ) - clone = Mock(return_value=(SimpleNamespace(id="draft-agent"), "draft-snapshot")) - monkeypatch.setattr(WorkflowAgentPublishService, "_clone_inline_graph_binding_for_node", clone) + sqlite_session.add_all([source_agent, source_snapshot, source]) + sqlite_session.commit() + target_agent = _inline_agent(agent_id="draft-agent", workflow_id="draft-workflow", node_id="agent-node") + target_snapshot = _snapshot(snapshot_id="draft-snapshot", agent_id=target_agent.id) + clone = Mock(return_value=(target_agent, target_snapshot)) + monkeypatch.setattr(AgentDslService, "clone_inline_binding_for_node", clone) WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( - session=session, + session=sqlite_session, source_workflow=_workflow(workflow_id="published-workflow", version="published"), draft_workflow=_workflow(workflow_id="draft-workflow"), account_id="account-2", ) clone.assert_called_once() - restored = session.add.call_args.args[0] + restored = sqlite_session.scalar( + select(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.workflow_id == "draft-workflow", + WorkflowAgentNodeBinding.workflow_version == Workflow.VERSION_DRAFT, + ) + ) + assert restored is not None assert restored.agent_id == "draft-agent" assert restored.current_snapshot_id == "draft-snapshot" 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 733c09896de..2a239c6e726 100644 --- a/api/tests/unit_tests/services/agent/test_workspace_service.py +++ b/api/tests/unit_tests/services/agent/test_workspace_service.py @@ -105,11 +105,6 @@ def test_workspace_client_honors_the_configured_snapshot_timeout(monkeypatch: py assert client._timeout == 123.5 -@pytest.mark.parametrize( - "sqlite_session", - [(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], - indirect=True, -) def test_create_binding_success_persists_new_workspace_and_binding( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: @@ -148,11 +143,6 @@ def test_create_binding_success_persists_new_workspace_and_binding( assert request.home_snapshot_ref == "home-ref" -@pytest.mark.parametrize( - "sqlite_session", - [(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], - indirect=True, -) def test_create_binding_without_home_snapshot_uses_backend_default( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: @@ -176,11 +166,6 @@ def test_create_binding_without_home_snapshot_uses_backend_default( assert request.home_snapshot_ref is None -@pytest.mark.parametrize( - "sqlite_session", - [(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], - indirect=True, -) def test_create_binding_rejects_missing_explicit_home_snapshot_before_backend_call( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: @@ -200,11 +185,6 @@ def test_create_binding_rejects_missing_explicit_home_snapshot_before_backend_ca client.create_execution_binding_sync.assert_not_called() -@pytest.mark.parametrize( - "sqlite_session", - [(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], - indirect=True, -) def test_create_second_binding_reuses_existing_workspace( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: @@ -243,7 +223,6 @@ def test_create_second_binding_reuses_existing_workspace( assert request.workspace_id == workspace.id -@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) def test_get_active_binding_resolves_exact_participant(sqlite_session: Session) -> None: conversation_workspace = _workspace(workspace_id="workspace-conversation") build_workspace = _workspace( @@ -274,7 +253,6 @@ def test_get_active_binding_resolves_exact_participant(sqlite_session: Session) assert resolved.id == conversation_binding.id -@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) def test_get_active_binding_rejects_wrong_owner(sqlite_session: Session) -> None: build_workspace = _workspace( workspace_id="workspace-build", @@ -299,7 +277,6 @@ def test_get_active_binding_rejects_wrong_owner(sqlite_session: Session) -> None assert resolved is None -@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) def test_retire_non_final_binding_keeps_workspace_active(sqlite_session: Session) -> None: binding = _binding() other_binding = _binding(binding_id="binding-2", agent_id="agent-2") @@ -320,7 +297,6 @@ def test_retire_non_final_binding_keeps_workspace_active(sqlite_session: Session assert other_binding.status is AgentWorkingResourceStatus.ACTIVE -@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) def test_retire_final_binding_retires_workspace(sqlite_session: Session) -> None: binding = _binding() workspace = _workspace() @@ -335,15 +311,14 @@ def test_retire_final_binding_retires_workspace(sqlite_session: Session) -> None assert workspace.retired_at == binding.retired_at -def test_retire_workspace_retires_all_active_bindings() -> None: +def test_retire_workspace_retires_all_active_bindings(sqlite_session: Session) -> None: workspace = _workspace() bindings = [_binding(), _binding(binding_id="binding-2", agent_id="agent-2")] - session = MagicMock() - session.scalar.return_value = workspace - session.scalars.return_value.all.return_value = bindings + sqlite_session.add_all([workspace, *bindings]) + sqlite_session.flush() retired_id = AgentWorkspaceService.retire_workspace( - session=session, + session=sqlite_session, tenant_id="tenant-1", workspace_id=workspace.id, ) @@ -354,7 +329,6 @@ def test_retire_workspace_retires_all_active_bindings() -> None: assert all(binding.retired_at == workspace.retired_at for binding in bindings) -@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) def test_retire_all_for_app_retires_only_active_workspaces_for_that_app(sqlite_session: Session) -> None: active = _workspace(workspace_id="workspace-active", owner_id="conversation-active") already_retired = _workspace( @@ -390,7 +364,6 @@ def test_retire_all_for_app_retires_only_active_workspaces_for_that_app(sqlite_s assert other_binding.status is AgentWorkingResourceStatus.ACTIVE -@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) def test_collect_binding_without_retired_workspace_destroys_binding_only( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: @@ -414,7 +387,6 @@ def test_collect_binding_without_retired_workspace_destroys_binding_only( assert sqlite_session.get(AgentWorkspace, workspace.id) is not None -@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) def test_collect_workspace_destroys_workspace_then_remaining_bindings( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: diff --git a/api/tests/unit_tests/services/test_app_generate_service.py b/api/tests/unit_tests/services/test_app_generate_service.py index b0bf1a2fd4e..613c6c2000e 100644 --- a/api/tests/unit_tests/services/test_app_generate_service.py +++ b/api/tests/unit_tests/services/test_app_generate_service.py @@ -21,6 +21,7 @@ from unittest.mock import MagicMock import pytest from pytest_mock import MockerFixture +from sqlalchemy.orm import Session import services.app_generate_service as ags_module from core.app.entities.app_invoke_entities import InvokeFrom @@ -79,6 +80,12 @@ def _make_user() -> MagicMock: return user +class _RealSessionTest: + @pytest.fixture(autouse=True) + def _bind_unbound_session(self, unbound_session: Session) -> None: + self.session = unbound_session + + def _make_workflow(*, workflow_id: str = "workflow-id", created_by: str = "owner-id") -> MagicMock: workflow = MagicMock() workflow.id = workflow_id @@ -251,7 +258,7 @@ class TestGetMaxActiveRequests: # --------------------------------------------------------------------------- # generate – every AppMode branch # --------------------------------------------------------------------------- -class TestGenerate: +class TestGenerate(_RealSessionTest): """Tests for AppGenerateService.generate covering each mode.""" @pytest.fixture(autouse=True) @@ -280,7 +287,7 @@ class TestGenerate: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) assert result == {"result": "ok"} gen_spy.assert_called_once() @@ -301,7 +308,7 @@ class TestGenerate: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) assert result == {"result": "agent"} gen_spy.assert_called_once() @@ -317,7 +324,7 @@ class TestGenerate: side_effect=lambda x: x, ) app = _make_app(AppMode.CHAT, is_agent=True) - session = MagicMock() + session = self.session result = AppGenerateService.generate( app_model=app, user=_make_user(), @@ -340,7 +347,7 @@ class TestGenerate: "services.app_generate_service.AgentAppGenerator.convert_to_event_stream", side_effect=lambda x: x, ) - session = MagicMock() + session = self.session result = AppGenerateService.generate( app_model=_make_app(AppMode.AGENT), @@ -371,7 +378,7 @@ class TestGenerate: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) assert result == {"result": "chat"} gen_spy.assert_called_once() @@ -391,7 +398,7 @@ class TestGenerate: side_effect=lambda x: x, ) - session = MagicMock() + session = self.session result = AppGenerateService.generate( app_model=_make_app(AppMode.ADVANCED_CHAT), user=_make_user(), @@ -430,7 +437,7 @@ class TestGenerate: args={"workflow_id": None, "query": "hi", "inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=True, - session=MagicMock(), + session=self.session, ) # In streaming mode it should go through retrieve_events, not generate gen_instance.retrieve_events.assert_called_once() @@ -453,7 +460,7 @@ class TestGenerate: side_effect=lambda x: x, ) - session = MagicMock() + session = self.session result = AppGenerateService.generate( app_model=_make_app(AppMode.WORKFLOW), user=_make_user(), @@ -492,7 +499,7 @@ class TestGenerate: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=True, - session=MagicMock(), + session=self.session, ) retrieve_spy.assert_called_once() # Dispatch is gated on subscribe; simulate the SSE layer entering the @@ -511,14 +518,14 @@ class TestGenerate: args={}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) # --------------------------------------------------------------------------- # generate – billing / quota # --------------------------------------------------------------------------- -class TestGenerateBilling: +class TestGenerateBilling(_RealSessionTest): @pytest.fixture(autouse=True) def _common(self, mocker: MockerFixture): mocker.patch("services.app_generate_service.RateLimit", _DummyRateLimit) @@ -549,7 +556,7 @@ class TestGenerateBilling: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) reserve_mock.assert_called_once_with(QuotaType.WORKFLOW, "tenant-id") quota_charge.commit.assert_called_once() @@ -573,7 +580,7 @@ class TestGenerateBilling: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) def test_exception_refunds_quota_and_exits_rate_limit( @@ -601,7 +608,7 @@ class TestGenerateBilling: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) quota_charge.refund.assert_called_once() @@ -633,7 +640,7 @@ class TestGenerateBilling: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) # exit is called in finally block for non-streaming assert exit_calls == ["dummy-request-id"] @@ -664,7 +671,7 @@ class TestGenerateBilling: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=MagicMock(), + session=self.session, ) quota_charge.refund.assert_called_once() @@ -698,7 +705,7 @@ class TestGenerateBilling: args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=True, - session=MagicMock(), + session=self.session, ) quota_charge.refund.assert_called_once() @@ -708,14 +715,16 @@ class TestGenerateBilling: # --------------------------------------------------------------------------- # _get_workflow # --------------------------------------------------------------------------- -class TestGetWorkflow: +class TestGetWorkflow(_RealSessionTest): def test_debugger_fetches_draft(self, mocker: MockerFixture): draft_wf = _make_workflow() ws = MagicMock() ws.get_draft_workflow.return_value = draft_wf mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - result = AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=MagicMock()) + result = AppGenerateService._get_workflow( + _make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=self.session + ) assert result is draft_wf ws.get_draft_workflow.assert_called_once() @@ -725,7 +734,7 @@ class TestGetWorkflow: mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) with pytest.raises(ValueError, match="Workflow not initialized"): - AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=MagicMock()) + AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=self.session) def test_non_debugger_fetches_published(self, mocker: MockerFixture): pub_wf = _make_workflow() @@ -734,7 +743,7 @@ class TestGetWorkflow: mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) result = AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=MagicMock() + _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=self.session ) assert result is pub_wf ws.get_published_workflow.assert_called_once() @@ -745,7 +754,7 @@ class TestGetWorkflow: mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) with pytest.raises(ValueError, match="Workflow not published"): - AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=MagicMock()) + AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=self.session) def test_specific_workflow_id_valid_uuid(self, mocker: MockerFixture): valid_uuid = str(uuid.uuid4()) @@ -758,7 +767,7 @@ class TestGetWorkflow: _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, workflow_id=valid_uuid, - session=MagicMock(), + session=self.session, ) assert result is specific_wf ws.get_published_workflow_by_id.assert_called_once() @@ -772,7 +781,7 @@ class TestGetWorkflow: _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, workflow_id="not-a-uuid", - session=MagicMock(), + session=self.session, ) def test_specific_workflow_id_not_found(self, mocker: MockerFixture): @@ -786,14 +795,14 @@ class TestGetWorkflow: _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, workflow_id=valid_uuid, - session=MagicMock(), + session=self.session, ) # --------------------------------------------------------------------------- # generate_single_iteration # --------------------------------------------------------------------------- -class TestGenerateSingleIteration: +class TestGenerateSingleIteration(_RealSessionTest): def test_advanced_chat_mode(self, mocker: MockerFixture): workflow = _make_workflow() mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) @@ -806,7 +815,7 @@ class TestGenerateSingleIteration: return_value={"event": "iteration"}, ) app = _make_app(AppMode.ADVANCED_CHAT) - session = MagicMock() + session = self.session result = AppGenerateService.generate_single_iteration( app_model=app, user=_make_user(), @@ -830,7 +839,7 @@ class TestGenerateSingleIteration: return_value={"event": "wf-iteration"}, ) app = _make_app(AppMode.WORKFLOW) - session = MagicMock() + session = self.session result = AppGenerateService.generate_single_iteration( app_model=app, user=_make_user(), @@ -846,14 +855,14 @@ class TestGenerateSingleIteration: app = _make_app(AppMode.CHAT) with pytest.raises(ValueError, match="Invalid app mode"): AppGenerateService.generate_single_iteration( - app_model=app, user=_make_user(), node_id="n1", args={}, session=MagicMock() + app_model=app, user=_make_user(), node_id="n1", args={}, session=self.session ) # --------------------------------------------------------------------------- # generate_single_loop # --------------------------------------------------------------------------- -class TestGenerateSingleLoop: +class TestGenerateSingleLoop(_RealSessionTest): def test_advanced_chat_mode(self, mocker: MockerFixture): workflow = _make_workflow() mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) @@ -866,7 +875,7 @@ class TestGenerateSingleLoop: return_value={"event": "loop"}, ) app = _make_app(AppMode.ADVANCED_CHAT) - session = MagicMock() + session = self.session result = AppGenerateService.generate_single_loop( app_model=app, user=_make_user(), @@ -890,7 +899,7 @@ class TestGenerateSingleLoop: return_value={"event": "wf-loop"}, ) app = _make_app(AppMode.WORKFLOW) - session = MagicMock() + session = self.session result = AppGenerateService.generate_single_loop( app_model=app, user=_make_user(), @@ -906,20 +915,20 @@ class TestGenerateSingleLoop: app = _make_app(AppMode.COMPLETION) with pytest.raises(ValueError, match="Invalid app mode"): AppGenerateService.generate_single_loop( - app_model=app, user=_make_user(), node_id="n1", args=MagicMock(), session=MagicMock() + app_model=app, user=_make_user(), node_id="n1", args=MagicMock(), session=self.session ) # --------------------------------------------------------------------------- # generate_more_like_this # --------------------------------------------------------------------------- -class TestGenerateMoreLikeThis: +class TestGenerateMoreLikeThis(_RealSessionTest): def test_delegates_to_completion_generator(self, mocker: MockerFixture): gen_spy = mocker.patch( "services.app_generate_service.CompletionAppGenerator.generate_more_like_this", return_value={"result": "similar"}, ) - session = MagicMock() + session = self.session result = AppGenerateService.generate_more_like_this( app_model=_make_app(AppMode.COMPLETION), user=_make_user(), From 280f81757d5fdea9c2dd804483c7f70c2e95cd9a Mon Sep 17 00:00:00 2001 From: Joel Date: Mon, 31 Aug 2026 06:58:56 +0000 Subject: [PATCH 15/21] fix: stabilize node block selectors and restore start tabs (#41523) --- .../workflow/__tests__/custom-edge.spec.tsx | 23 ++++++++++++++----- .../components/__tests__/node-handle.spec.tsx | 21 +++++++++++++++++ .../nodes/_base/components/node-handle.tsx | 4 ++++ 3 files changed, 42 insertions(+), 6 deletions(-) diff --git a/web/app/components/workflow/__tests__/custom-edge.spec.tsx b/web/app/components/workflow/__tests__/custom-edge.spec.tsx index d7a61199c20..77bd5373d67 100644 --- a/web/app/components/workflow/__tests__/custom-edge.spec.tsx +++ b/web/app/components/workflow/__tests__/custom-edge.spec.tsx @@ -1,9 +1,11 @@ import type { ReactNode } from 'react' -import { render, screen } from '@testing-library/react' +import { screen } from '@testing-library/react' +import userEvent from '@testing-library/user-event' import { Position } from 'reactflow' import { ErrorHandleTypeEnum } from '@/app/components/workflow/nodes/_base/components/error-handle/types' import CustomEdge from '../custom-edge' import { BlockEnum, NodeRunningStatus } from '../types' +import { renderWorkflowComponent } from './workflow-test-env' const mockUseAvailableBlocks = vi.hoisted(() => vi.fn()) const mockUseNodesInteractions = vi.hoisted(() => vi.fn()) @@ -38,6 +40,9 @@ vi.mock('reactflow', () => ({ Right: 'right', Left: 'left', }, + useStoreApi: () => ({ + getState: () => ({ getNodes: () => [] }), + }), })) vi.mock('../hooks/use-available-blocks', async (importOriginal) => { @@ -81,8 +86,10 @@ describe('CustomEdge', () => { }) }) - it('should render a gradient edge and its real insert-node trigger', () => { - render( + it('should render a gradient edge and hide the start tab from its insert-node selector', async () => { + const user = userEvent.setup() + + renderWorkflowComponent( { opacity: '0.7', zIndex: '1001', }) + + await user.click(addBlockTrigger) + + expect(screen.queryByRole('tab', { name: 'workflow.tabs.start' })).not.toBeInTheDocument() }) it('should prefer the running stroke color when the edge is selected', () => { - render( + renderWorkflowComponent( { }) it('should use the fail-branch running color while the connected node is hovering', () => { - render( + renderWorkflowComponent( { }) it('should fall back to the default edge color when no highlight state is active', () => { - render( + renderWorkflowComponent( { // Target-side tests cover selector visibility, connection locking, and status rendering. describe('NodeTargetHandle', () => { + it('should show the start tab when adding a node before the target node', async () => { + const user = userEvent.setup() + + renderTargetHandle() + + await user.click(screen.getByTestId('handle-target-handle')) + + expect(screen.getByRole('tab', { name: 'workflow.tabs.start' })).toBeInTheDocument() + }) + it('should toggle the target add trigger', () => { renderTargetHandle() @@ -260,6 +271,16 @@ describe('node-handle', () => { // Source-side tests cover selector opening paths, previous-node selection, and status styling. describe('NodeSourceHandle', () => { + it('should show the start tab when adding a node after the source node', async () => { + const user = userEvent.setup() + + renderSourceHandle() + + await user.click(screen.getByTestId('handle-source-handle')) + + expect(screen.getByRole('tab', { name: 'workflow.tabs.start' })).toBeInTheDocument() + }) + it('should toggle the source add trigger', () => { renderSourceHandle() diff --git a/web/app/components/workflow/nodes/_base/components/node-handle.tsx b/web/app/components/workflow/nodes/_base/components/node-handle.tsx index 955e5eb2f87..30a22ab512c 100644 --- a/web/app/components/workflow/nodes/_base/components/node-handle.tsx +++ b/web/app/components/workflow/nodes/_base/components/node-handle.tsx @@ -79,6 +79,7 @@ export const NodeTargetHandle = memo( 'z-1 size-4! rounded-none! border-none! bg-transparent! outline-hidden!', 'after:absolute after:top-1 after:left-1.5 after:h-2 after:w-0.5 after:bg-workflow-link-line-handle', 'transition-all hover:scale-125', + open && 'scale-125', data._runningStatus === NodeRunningStatus.Succeeded && 'after:bg-workflow-link-line-success-handle', data._runningStatus === NodeRunningStatus.Failed && @@ -106,6 +107,7 @@ export const NodeTargetHandle = memo( nextNodeTargetHandle: handleId, }} placement="left" + showStartTab triggerClassName={` absolute left-0 top-0 opacity-0 pointer-events-none transition-opacity duration-150 ${nodeSelectorClassName} @@ -206,6 +208,7 @@ export const NodeSourceHandle = memo( 'group/handle z-1 size-4! rounded-none! border-none! bg-transparent! outline-hidden!', 'after:absolute after:top-1 after:right-1.5 after:h-2 after:w-0.5 after:bg-workflow-link-line-handle', 'transition-all hover:scale-125', + open && 'scale-125', data._runningStatus === NodeRunningStatus.Succeeded && 'after:bg-workflow-link-line-success-handle', data._runningStatus === NodeRunningStatus.Failed && @@ -252,6 +255,7 @@ export const NodeSourceHandle = memo( data-popup-open:opacity-100 `} availableBlocksTypes={availableNextBlocks} + showStartTab /> )} From f021ee4dc0ced54c7cb461fdcf50cea7d6047602 Mon Sep 17 00:00:00 2001 From: Joel Date: Mon, 31 Aug 2026 07:23:45 +0000 Subject: [PATCH 16/21] fix: prevent the exit versions button from shrinking (#41527) --- .../configure/components/orchestrate/publish-bar/index.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/features/agent-v2/agent-detail/configure/components/orchestrate/publish-bar/index.tsx b/web/features/agent-v2/agent-detail/configure/components/orchestrate/publish-bar/index.tsx index d1fcfb76586..a04a731b2e8 100644 --- a/web/features/agent-v2/agent-detail/configure/components/orchestrate/publish-bar/index.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/orchestrate/publish-bar/index.tsx @@ -471,7 +471,7 @@ function AgentVersionRestoreBar({