test: migrate agent app sandbox sessions and ORM models to SQLite (#40568)

api/tests/unit_tests/services/test_agent_app_sandbox_service.py

@@ -8,7 +8,7 @@ from unittest.mock import MagicMock
 import pytest
 from dify_agent.client import Client
 from dify_agent.protocol import BindingFileDownloadResponse, BindingFileListResponse, BindingFileReadResponse
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 from graphon.enums import WorkflowNodeExecutionStatus
 from models.agent import (
@@ -30,7 +30,6 @@ from models.enums import ConversationFromSource, CreatorUserRole
 from models.model import App, AppMode, Conversation, IconType
 from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom
 from services import agent_app_sandbox_service as sandbox_module
-from services.agent.workspace_service import AgentWorkspaceService
 from services.agent_app_sandbox_service import (
     AgentAppSandboxService,
     AgentSandboxDownload,
@@ -149,6 +148,8 @@ def _add_binding(
     owner_id: str,
     owner_scope_key: str = "root",
     status: AgentWorkingResourceStatus = AgentWorkingResourceStatus.ACTIVE,
+    agent_config_version_id: str | None = None,
+    agent_config_version_kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT,
 ) -> AgentWorkspaceBinding:
     active_guard = 1 if status is AgentWorkingResourceStatus.ACTIVE else None
     workspace = AgentWorkspace(
@@ -169,8 +170,8 @@ def _add_binding(
         workspace_id=workspace_id,
         agent_id=agent_id,
         base_home_snapshot_id=None,
-        agent_config_version_id=f"{binding_id}-config",
-        agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT,
+        agent_config_version_id=agent_config_version_id or f"{binding_id}-config",
+        agent_config_version_kind=agent_config_version_kind,
         backend_binding_ref=f"{binding_id}-ref",
         status=status,
     )
@@ -178,6 +179,52 @@ def _add_binding(
     return binding
 
 
+def _add_build_draft_caller(
+    session: Session,
+    *,
+    parent_app_id: str = "app-1",
+    backing_app_id: str | None = None,
+    runtime_app_id: str = "app-1",
+) -> AgentWorkspaceBinding:
+    session.add(
+        Agent(
+            id="agent-1",
+            tenant_id="tenant-1",
+            name="Agent",
+            description="",
+            agent_kind=AgentKind.DIFY_AGENT,
+            scope=AgentScope.WORKFLOW_ONLY if backing_app_id else AgentScope.ROSTER,
+            source=AgentSource.WORKFLOW if backing_app_id else AgentSource.AGENT_APP,
+            app_id=parent_app_id,
+            backing_app_id=backing_app_id,
+            status=AgentStatus.ACTIVE,
+        )
+    )
+    binding = _add_binding(
+        session,
+        binding_id="binding-build",
+        workspace_id="workspace-build",
+        app_id=runtime_app_id,
+        owner_type=AgentWorkspaceOwnerType.BUILD_DRAFT,
+        owner_id="build-1",
+        agent_config_version_id="config-1",
+        agent_config_version_kind=AgentConfigVersionKind.DRAFT,
+    )
+    session.add(
+        AgentConfigDraft(
+            id="build-1",
+            tenant_id="tenant-1",
+            agent_id="agent-1",
+            draft_type=AgentConfigDraftType.DEBUG_BUILD,
+            account_id="account-1",
+            draft_owner_key="account-1",
+            agent_workspace_binding_id=binding.id,
+            config_snapshot=AgentSoulConfig(),
+        )
+    )
+    return binding
+
+
 def _download_client() -> MagicMock:
     client = MagicMock()
     client.download_binding_file_sync.return_value = BindingFileDownloadResponse(reference="dify-file-ref:canonical")
@@ -545,7 +592,8 @@ def _workflow_execution(
     app_id: str = "app-1",
     workflow_run_id: str = "run-1",
     node_id: str = "node-1",
-    binding_id: str,
+    binding_id: str | None,
+    workflow_agent_binding_id: str | None = "workflow-binding-1",
     created_by: str = "historical-account",
 ) -> WorkflowNodeExecutionModel:
     return WorkflowNodeExecutionModel(
@@ -561,7 +609,9 @@ def _workflow_execution(
         title=node_id,
         agent_workspace_binding_id=binding_id,
         inputs=None,
-        process_data=json.dumps({"workflow_agent_binding_id": "workflow-binding-1"}),
+        process_data=json.dumps(
+            {"workflow_agent_binding_id": workflow_agent_binding_id} if workflow_agent_binding_id is not None else {}
+        ),
         outputs=None,
         status=WorkflowNodeExecutionStatus.SUCCEEDED,
         error=None,
@@ -696,24 +746,19 @@ def test_workflow_download_resolves_only_exact_active_owner_chain(
 )
 def test_agent_app_file_browsing_uses_build_draft_caller(
     monkeypatch: pytest.MonkeyPatch,
+    sqlite_session: Session,
     parent_app_id: str,
     backing_app_id: str | None,
     runtime_app_id: str,
 ) -> None:
-    session = MagicMock()
-    session.scalar.side_effect = [
-        SimpleNamespace(app_id=parent_app_id, backing_app_id=backing_app_id),
-        SimpleNamespace(agent_workspace_binding_id="binding-build"),
-    ]
-    _use_session(monkeypatch, session)
-    binding = SimpleNamespace(
-        agent_id="agent-1",
-        backend_binding_ref="binding-build-ref",
-        agent_config_version_id="config-1",
-        agent_config_version_kind=AgentConfigVersionKind.DRAFT,
+    _add_build_draft_caller(
+        sqlite_session,
+        parent_app_id=parent_app_id,
+        backing_app_id=backing_app_id,
+        runtime_app_id=runtime_app_id,
     )
-    get_binding = MagicMock(return_value=binding)
-    monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding)
+    sqlite_session.commit()
+    _use_session(monkeypatch, sqlite_session)
     client = MagicMock()
     response = BindingFileListResponse(path=".", entries=[], truncated=False)
     client.list_binding_files_sync.return_value = response
@@ -729,28 +774,22 @@ def test_agent_app_file_browsing_uses_build_draft_caller(
     )
 
     assert result is response
-    owner_scope = get_binding.call_args.kwargs["expected_owner_scope"]
-    assert owner_scope.app_id == runtime_app_id
-    assert owner_scope.owner_type is AgentWorkspaceOwnerType.BUILD_DRAFT
-    assert owner_scope.owner_id == "build-1"
     client.list_binding_files_sync.assert_called_once_with("binding-build-ref", ".")
 
 
-def test_workflow_file_access_uses_node_execution_pointer(monkeypatch: pytest.MonkeyPatch) -> None:
-    execution = SimpleNamespace(
-        agent_workspace_binding_id="binding-workflow",
-        process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"},
+def test_workflow_file_access_uses_node_execution_pointer(
+    sqlite_session: Session,
+) -> None:
+    binding = _add_binding(
+        sqlite_session,
+        binding_id="binding-workflow",
+        workspace_id="workspace-workflow",
+        owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN,
+        owner_id="run-1",
+        owner_scope_key="node-1:workflow-binding-1",
     )
-    session = MagicMock()
-    session.scalar.return_value = execution
-    binding = SimpleNamespace(
-        agent_id="agent-1",
-        backend_binding_ref="binding-workflow-ref",
-        agent_config_version_id="config-1",
-        agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT,
-    )
-    get_binding = MagicMock(return_value=binding)
-    monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding)
+    sqlite_session.add(_workflow_execution(execution_id="execution-1", binding_id=binding.id))
+    sqlite_session.commit()
     client = MagicMock()
     response = BindingFileReadResponse(path="report.txt", size=2, truncated=False, binary=False, text="ok")
     client.read_binding_file_sync.return_value = response
@@ -762,39 +801,37 @@ def test_workflow_file_access_uses_node_execution_pointer(monkeypatch: pytest.Mo
         node_id="node-1",
         node_execution_id="execution-1",
         path="report.txt",
-        session=session,
+        session=sqlite_session,
     )
 
     assert result is response
-    assert get_binding.call_args.kwargs["binding_id"] == "binding-workflow"
     client.read_binding_file_sync.assert_called_once_with("binding-workflow-ref", "report.txt")
-    session.rollback.assert_called_once_with()
+    assert not sqlite_session.in_transaction()
 
 
 def test_workflow_download_uses_authenticated_account_and_trusted_file_request(
     monkeypatch: pytest.MonkeyPatch,
+    sqlite_session: Session,
+    sqlite_session_factory: sessionmaker[Session],
 ) -> None:
-    execution = SimpleNamespace(
-        agent_workspace_binding_id="binding-workflow",
-        process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"},
+    binding = _add_binding(
+        sqlite_session,
+        binding_id="binding-workflow",
+        workspace_id="workspace-workflow",
+        owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN,
+        owner_id="run-1",
+        owner_scope_key="node-1:workflow-binding-1",
     )
-    session = MagicMock()
-    session.scalar.return_value = execution
-    binding = SimpleNamespace(
-        agent_id="agent-1",
-        backend_binding_ref="binding-workflow-ref",
-        agent_config_version_id="config-1",
-        agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT,
-    )
-    monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding))
+    sqlite_session.add(_workflow_execution(execution_id="execution-1", binding_id=binding.id))
+    sqlite_session.commit()
     events: list[str] = []
 
     @contextmanager
     def session_scope():
-        try:
-            yield session
-        finally:
-            events.append("session-exit")
+        with sqlite_session_factory() as service_session:
+            yield service_session
+            assert not service_session.in_transaction()
+        events.append("session-exit")
 
     monkeypatch.setattr(sandbox_module.session_factory, "create_session", session_scope)
     client = MagicMock()
@@ -826,7 +863,6 @@ def test_workflow_download_uses_authenticated_account_and_trusted_file_request(
     assert request.execution_context.user_id == "account-1"
     assert request.execution_context.user_from == "account"
     assert request.execution_context.node_execution_id == "execution-1"
-    session.rollback.assert_called_once_with()
     assert events == ["session-exit", "client-download", "file-request"]
     request_download.assert_called_once_with(
         tenant_id="tenant-1",
@@ -840,29 +876,20 @@ def test_workflow_download_uses_authenticated_account_and_trusted_file_request(
 
 def test_agent_app_download_uses_complete_account_context_after_session_exit(
     monkeypatch: pytest.MonkeyPatch,
+    sqlite_session: Session,
+    sqlite_session_factory: sessionmaker[Session],
 ) -> None:
     events: list[str] = []
-    session = MagicMock()
-    session.scalar.side_effect = [
-        SimpleNamespace(app_id="app-1", backing_app_id=None),
-        SimpleNamespace(agent_workspace_binding_id="binding-build"),
-    ]
+    _add_build_draft_caller(sqlite_session)
+    sqlite_session.commit()
 
     @contextmanager
     def session_scope():
-        try:
-            yield session
-        finally:
-            events.append("session-exit")
+        with sqlite_session_factory() as service_session:
+            yield service_session
+        events.append("session-exit")
 
     monkeypatch.setattr(sandbox_module.session_factory, "create_session", session_scope)
-    binding = SimpleNamespace(
-        agent_id="agent-1",
-        backend_binding_ref="binding-build-ref",
-        agent_config_version_id="config-1",
-        agent_config_version_kind=AgentConfigVersionKind.DRAFT,
-    )
-    monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding))
     client = MagicMock()
     client.download_binding_file_sync.side_effect = lambda _request: (
         events.append("client-download") or BindingFileDownloadResponse(reference="dify-file-ref:canonical")
@@ -934,31 +961,27 @@ def test_file_request_rejection_maps_to_download_unavailable(monkeypatch: pytest
 
 
 @pytest.mark.parametrize(
-    "execution",
+    ("binding_id", "workflow_agent_binding_id"),
     [
-        pytest.param(
-            SimpleNamespace(
-                agent_workspace_binding_id=None,
-                process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"},
-            ),
-            id="missing-binding-pointer",
-        ),
-        pytest.param(
-            SimpleNamespace(
-                agent_workspace_binding_id="binding-workflow",
-                process_data_dict={},
-            ),
-            id="missing-process-data",
-        ),
+        pytest.param(None, "workflow-binding-1", id="missing-binding-pointer"),
+        pytest.param("binding-workflow", None, id="missing-process-data"),
     ],
 )
 def test_workflow_download_rejects_missing_binding_metadata_before_network(
     monkeypatch: pytest.MonkeyPatch,
-    execution: SimpleNamespace,
+    sqlite_session: Session,
+    binding_id: str | None,
+    workflow_agent_binding_id: str | None,
 ) -> None:
-    session = MagicMock()
-    session.scalar.return_value = execution
-    _use_session(monkeypatch, session)
+    sqlite_session.add(
+        _workflow_execution(
+            execution_id="execution-1",
+            binding_id=binding_id,
+            workflow_agent_binding_id=workflow_agent_binding_id,
+        )
+    )
+    sqlite_session.commit()
+    _use_session(monkeypatch, sqlite_session)
     client = MagicMock()
 
     with pytest.raises(AgentSandboxInspectorError) as exc_info:
@@ -978,15 +1001,19 @@ def test_workflow_download_rejects_missing_binding_metadata_before_network(
 
 def test_workflow_download_rejects_non_active_or_mismatched_binding_before_network(
     monkeypatch: pytest.MonkeyPatch,
+    sqlite_session: Session,
 ) -> None:
-    execution = SimpleNamespace(
-        agent_workspace_binding_id="binding-workflow",
-        process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"},
+    binding = _add_binding(
+        sqlite_session,
+        binding_id="binding-workflow",
+        workspace_id="workspace-workflow",
+        owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN,
+        owner_id="run-other",
+        owner_scope_key="node-1:workflow-binding-1",
     )
-    session = MagicMock()
-    session.scalar.return_value = execution
-    _use_session(monkeypatch, session)
-    monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None))
+    sqlite_session.add(_workflow_execution(execution_id="execution-1", binding_id=binding.id))
+    sqlite_session.commit()
+    _use_session(monkeypatch, sqlite_session)
     client = MagicMock()
 
     with pytest.raises(AgentSandboxInspectorError) as exc_info: