commit 318c2c2edae3876fcf8ba9b4a2b208206631cfa9 Author: upstream Date: Tue Sep 29 11:33:44 2026 +0200 Upstream langgenius/dify state of the touched files before 0ff8f0cd14 diff --git a/api/core/workflow/nodes/agent_v2/session_store.py b/api/core/workflow/nodes/agent_v2/session_store.py new file mode 100644 index 0000000..dd5f511 --- /dev/null +++ b/api/core/workflow/nodes/agent_v2/session_store.py @@ -0,0 +1,383 @@ +"""Workflow Agent participant persistence keyed by node execution.""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass + +from agenton.compositor import CompositorSessionSnapshot +from sqlalchemy import select +from sqlalchemy.orm import Session + +from core.db.session_factory import session_factory +from models.agent import ( + AgentConfigVersionKind, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) +from models.workflow import WorkflowNodeExecutionModel +from services.agent.workspace_service import ( + AgentWorkspaceNotFoundError, + AgentWorkspaceService, + WorkspaceOwnerScope, +) + +_CALLER_VISIBILITY_ATTEMPTS = 60 +_CALLER_VISIBILITY_INTERVAL_SECONDS = 0.05 + + +def resolve_workflow_agent_workspace_owner_scope( + *, + tenant_id: str, + app_id: str, + conversation_id: str | None, + workflow_run_id: str | None, + node_id: str, + workflow_agent_binding_id: str, + node_execution_id: str | None = None, +) -> WorkspaceOwnerScope: + """Choose the Workspace owner for a workflow Agent participant. + + Chatflow runs carry ``conversation_id`` and should persist Agent memory across + turns the same way standalone Agent Apps do. Pure workflow runs keep the + per-run ``WORKFLOW_RUN`` scope. + """ + + owner_scope_key = f"{node_id}:{workflow_agent_binding_id}" + if conversation_id: + return WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=conversation_id, + owner_scope_key=owner_scope_key, + ) + return WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN, + owner_id=workflow_run_id or node_execution_id or "", + owner_scope_key=owner_scope_key, + ) + + +@dataclass(frozen=True, slots=True) +class WorkflowAgentSessionScope: + tenant_id: str + app_id: str + workflow_id: str + workflow_run_id: str | None + node_id: str + node_execution_id: str + workflow_agent_binding_id: str + agent_id: str + agent_config_snapshot_id: str + conversation_id: str | None = None + + @property + def workspace_owner(self) -> WorkspaceOwnerScope: + return resolve_workflow_agent_workspace_owner_scope( + tenant_id=self.tenant_id, + app_id=self.app_id, + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + node_id=self.node_id, + workflow_agent_binding_id=self.workflow_agent_binding_id, + node_execution_id=self.node_execution_id, + ) + + +@dataclass(frozen=True, slots=True) +class StoredWorkflowAgentSession: + scope: WorkflowAgentSessionScope + binding_id: str + workspace_id: str + backend_binding_ref: str + session_snapshot: CompositorSessionSnapshot | None + pending_form_id: str | None = None + pending_tool_call_id: str | None = None + + +class WorkflowAgentWorkspaceStore: + """Load or create the participant named by a node execution caller row.""" + + @staticmethod + def load_active_participant( + *, + session: Session, + scope: WorkspaceOwnerScope, + agent_id: str, + ) -> AgentWorkspaceBinding | None: + """Look up the participant that pins a Chatflow node's generation.""" + + return AgentWorkspaceService.resolve_active_binding_for_scope( + session=session, + scope=scope, + agent_id=agent_id, + ) + + def load_existing_node_execution_scope( + self, + *, + tenant_id: str, + app_id: str, + workflow_id: str, + workflow_run_id: str | None, + node_id: str, + node_execution_id: str, + conversation_id: str | None = None, + ) -> WorkflowAgentSessionScope | None: + """Return the generation pinned by an existing node execution participant.""" + + with session_factory.create_session() as session: + execution = self._load_execution_by_identity( + session=session, + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run_id, + node_id=node_id, + node_execution_id=node_execution_id, + ) + binding_id = execution.agent_workspace_binding_id + if binding_id is None: + return None + process_data = execution.process_data_dict + if not isinstance(process_data, dict): + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is invalid") + workflow_agent_binding_id = process_data.get("workflow_agent_binding_id") + if not isinstance(workflow_agent_binding_id, str): + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is missing") + owner_scope = resolve_workflow_agent_workspace_owner_scope( + tenant_id=tenant_id, + app_id=app_id, + conversation_id=conversation_id, + workflow_run_id=workflow_run_id, + node_id=node_id, + workflow_agent_binding_id=workflow_agent_binding_id, + node_execution_id=node_execution_id, + ) + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=tenant_id, + binding_id=binding_id, + expected_owner_scope=owner_scope, + ) + if binding is None or binding.agent_config_version_kind != AgentConfigVersionKind.SNAPSHOT: + raise AgentWorkspaceNotFoundError("Workflow node participant Binding is unavailable") + return WorkflowAgentSessionScope( + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run_id, + node_id=node_id, + node_execution_id=node_execution_id, + workflow_agent_binding_id=workflow_agent_binding_id, + agent_id=binding.agent_id, + agent_config_snapshot_id=binding.agent_config_version_id, + conversation_id=conversation_id, + ) + + def load_or_create_node_execution_session( + self, scope: WorkflowAgentSessionScope, *, home_snapshot_id: str | None + ) -> StoredWorkflowAgentSession: + with session_factory.create_session() as session: + execution = self._load_execution(session=session, scope=scope) + process_data = execution.process_data_dict + if process_data is None: + process_data = {} + if not isinstance(process_data, dict): + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is invalid") + stored_workflow_binding_id = process_data.get("workflow_agent_binding_id") + if stored_workflow_binding_id is not None and stored_workflow_binding_id != scope.workflow_agent_binding_id: + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity does not match") + + binding_id = execution.agent_workspace_binding_id + if binding_id is None: + binding = self._resolve_or_create_binding( + session=session, + scope=scope, + home_snapshot_id=home_snapshot_id, + ) + execution.agent_workspace_binding_id = binding.id + execution.process_data = json.dumps( + { + **process_data, + "workflow_agent_binding_id": scope.workflow_agent_binding_id, + }, + ensure_ascii=False, + ) + session.commit() + else: + if stored_workflow_binding_id is None: + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is missing") + resolved_binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=scope.tenant_id, + binding_id=binding_id, + expected_owner_scope=scope.workspace_owner, + ) + if resolved_binding is None or resolved_binding.agent_id != scope.agent_id: + raise AgentWorkspaceNotFoundError("Workflow node participant Binding is unavailable") + binding = resolved_binding + AgentWorkspaceService.validate_binding_generation( + binding, + base_home_snapshot_id=home_snapshot_id, + agent_config_version_id=scope.agent_config_snapshot_id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + return self._stored(session, scope, binding) + + def save_active_snapshot( + self, + *, + scope: WorkflowAgentSessionScope, + binding_id: str, + snapshot: CompositorSessionSnapshot | None, + pending_form_id: str | None = None, + pending_tool_call_id: str | None = None, + ) -> None: + if snapshot is None: + return + AgentWorkspaceService.save_binding_session_snapshot( + tenant_id=scope.tenant_id, + binding_id=binding_id, + session_snapshot=snapshot.model_dump_json(), + pending_form_id=pending_form_id, + pending_tool_call_id=pending_tool_call_id, + ) + + def retire_workflow_run(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> list[str]: + """Retire active Workspaces, commit, and return active or already-retired IDs for collection.""" + + retired: list[str] = [] + with session_factory.create_session() as session: + workspaces = session.scalars( + select(AgentWorkspace).where( + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.app_id == app_id, + AgentWorkspace.owner_type == AgentWorkspaceOwnerType.WORKFLOW_RUN, + AgentWorkspace.owner_id == workflow_run_id, + AgentWorkspace.status.in_((AgentWorkingResourceStatus.ACTIVE, AgentWorkingResourceStatus.RETIRED)), + ) + ).all() + for workspace in workspaces: + if workspace.status == AgentWorkingResourceStatus.RETIRED: + retired.append(workspace.id) + continue + workspace_id = AgentWorkspaceService.retire_workspace( + session=session, + tenant_id=tenant_id, + workspace_id=workspace.id, + ) + if workspace_id is not None: + retired.append(workspace_id) + session.commit() + return retired + + @staticmethod + def _load_execution(*, session: Session, scope: WorkflowAgentSessionScope) -> WorkflowNodeExecutionModel: + return WorkflowAgentWorkspaceStore._load_execution_by_identity( + session=session, + tenant_id=scope.tenant_id, + app_id=scope.app_id, + workflow_id=scope.workflow_id, + workflow_run_id=scope.workflow_run_id, + node_id=scope.node_id, + node_execution_id=scope.node_execution_id, + ) + + @staticmethod + def _load_execution_by_identity( + *, + session: Session, + tenant_id: str, + app_id: str, + workflow_id: str, + workflow_run_id: str | None, + node_id: str, + node_execution_id: str, + ) -> WorkflowNodeExecutionModel: + """Wait briefly for the already-emitted node-start event to persist its caller row.""" + + stmt = select(WorkflowNodeExecutionModel).where( + WorkflowNodeExecutionModel.id == node_execution_id, + WorkflowNodeExecutionModel.tenant_id == tenant_id, + WorkflowNodeExecutionModel.app_id == app_id, + WorkflowNodeExecutionModel.workflow_id == workflow_id, + WorkflowNodeExecutionModel.node_id == node_id, + WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id, + ) + for attempt in range(_CALLER_VISIBILITY_ATTEMPTS): + execution = session.scalar(stmt) + if execution is not None: + return execution + if attempt < _CALLER_VISIBILITY_ATTEMPTS - 1: + time.sleep(_CALLER_VISIBILITY_INTERVAL_SECONDS) + + raise AgentWorkspaceNotFoundError("Workflow node execution caller is unavailable") + + @staticmethod + def _resolve_or_create_binding( + *, + session: Session, + scope: WorkflowAgentSessionScope, + home_snapshot_id: str | None, + ) -> AgentWorkspaceBinding: + """Reuse a conversation-scoped participant or allocate a new one.""" + + if scope.conversation_id is not None: + existing_binding = WorkflowAgentWorkspaceStore.load_active_participant( + session=session, + scope=scope.workspace_owner, + agent_id=scope.agent_id, + ) + if existing_binding is not None: + AgentWorkspaceService.validate_binding_generation( + existing_binding, + base_home_snapshot_id=home_snapshot_id, + agent_config_version_id=scope.agent_config_snapshot_id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + return existing_binding + + return AgentWorkspaceService.create_binding( + session=session, + scope=scope.workspace_owner, + agent_id=scope.agent_id, + base_home_snapshot_id=home_snapshot_id, + agent_config_version_id=scope.agent_config_snapshot_id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + + @staticmethod + def _stored( + session: Session, + scope: WorkflowAgentSessionScope, + binding: AgentWorkspaceBinding, + ) -> StoredWorkflowAgentSession: + snapshot = ( + CompositorSessionSnapshot.model_validate_json(binding.session_snapshot) + if binding.session_snapshot + else None + ) + return StoredWorkflowAgentSession( + scope=scope, + binding_id=binding.id, + workspace_id=binding.workspace_id, + backend_binding_ref=binding.backend_binding_ref, + session_snapshot=snapshot, + pending_form_id=binding.pending_form_id, + pending_tool_call_id=binding.pending_tool_call_id, + ) + + +__all__ = [ + "StoredWorkflowAgentSession", + "WorkflowAgentSessionScope", + "WorkflowAgentWorkspaceStore", + "resolve_workflow_agent_workspace_owner_scope", +] diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_store.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_store.py new file mode 100644 index 0000000..e18e36a --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_store.py @@ -0,0 +1,553 @@ +import json +from collections.abc import Iterator +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from agenton.compositor import CompositorSessionSnapshot +from agenton.compositor.schemas import LayerSessionSnapshot +from agenton.layers.base import LifecycleState +from sqlalchemy import Engine, event, func, select +from sqlalchemy.orm import Session, sessionmaker + +from core.workflow.nodes.agent_v2.session_store import ( + WorkflowAgentSessionScope, + WorkflowAgentWorkspaceStore, + resolve_workflow_agent_workspace_owner_scope, +) +from graphon.enums import WorkflowNodeExecutionStatus +from models.agent import ( + AgentConfigVersionKind, + AgentHomeSnapshot, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) +from models.enums import CreatorUserRole +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom +from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService + + +def _scope() -> WorkflowAgentSessionScope: + return WorkflowAgentSessionScope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + workflow_agent_binding_id="workflow-binding-1", + agent_id="agent-1", + agent_config_snapshot_id="config-1", + ) + + +def _execution_row( + *, + execution_id: str = "execution-1", + workflow_run_id: str = "run-1", + binding_id: str | None = None, + process_data: dict[str, object] | None = None, +) -> WorkflowNodeExecutionModel: + return WorkflowNodeExecutionModel( + id=execution_id, + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=workflow_run_id, + index=1, + predecessor_node_id=None, + node_execution_id="node-execution-1", + node_id="node-1", + node_type="agent", + title="Agent", + agent_workspace_binding_id=binding_id, + inputs="{}", + process_data=json.dumps(process_data) if process_data is not None else None, + outputs=None, + status=WorkflowNodeExecutionStatus.RUNNING, + error=None, + elapsed_time=0, + execution_metadata=None, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + finished_at=None, + ) + + +def _workspace_row( + *, + workspace_id: str = "workspace-1", + tenant_id: str = "tenant-1", + app_id: str = "app-1", + owner_scope_key: str = "node-1:workflow-binding-1", + status: AgentWorkingResourceStatus = AgentWorkingResourceStatus.ACTIVE, +) -> AgentWorkspace: + return AgentWorkspace( + id=workspace_id, + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN, + owner_id="run-1", + owner_scope_key=owner_scope_key, + backend_workspace_ref=f"{workspace_id}-ref", + status=status, + active_guard=1 if status is AgentWorkingResourceStatus.ACTIVE else None, + ) + + +def _binding_row( + *, + binding_id: str = "binding-1", + workspace_id: str = "workspace-1", + status: AgentWorkingResourceStatus = AgentWorkingResourceStatus.ACTIVE, + base_home_snapshot_id: str | None = "home-1", +) -> AgentWorkspaceBinding: + return AgentWorkspaceBinding( + id=binding_id, + tenant_id="tenant-1", + app_id="app-1", + workspace_id=workspace_id, + agent_id="agent-1", + base_home_snapshot_id=base_home_snapshot_id, + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="backend-binding-1" if binding_id == "binding-1" else f"{binding_id}-ref", + status=status, + ) + + +def _home_snapshot() -> AgentHomeSnapshot: + return AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + + +def _install_backend_client(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + client = MagicMock() + client.create_execution_binding_sync.return_value = SimpleNamespace( + binding_ref="backend-binding-1", + workspace_ref="workspace-1-ref", + ) + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + return client + + +@pytest.fixture +def executed_statements(sqlite_engine: Engine) -> Iterator[list[str]]: + statements: list[str] = [] + + def record_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: + statements.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", record_statement) + try: + yield statements + finally: + event.remove(sqlite_engine, "before_cursor_execute", record_statement) + + +def _execution_selects(statements: list[str]) -> list[str]: + return [statement for statement in statements if "FROM workflow_node_executions" in statement] + + +def test_scope_uses_node_and_workflow_binding_as_workspace_subscope() -> None: + owner = _scope().workspace_owner + assert owner.owner_type is AgentWorkspaceOwnerType.WORKFLOW_RUN + assert owner.owner_id == "run-1" + assert owner.owner_scope_key == "node-1:workflow-binding-1" + + +def test_resolve_workspace_owner_scope_uses_conversation_for_chatflow() -> None: + owner = resolve_workflow_agent_workspace_owner_scope( + tenant_id="tenant-1", + app_id="app-1", + conversation_id="conversation-1", + workflow_run_id="run-1", + node_id="node-1", + workflow_agent_binding_id="workflow-binding-1", + ) + assert owner.owner_type is AgentWorkspaceOwnerType.CONVERSATION + assert owner.owner_id == "conversation-1" + assert owner.owner_scope_key == "node-1:workflow-binding-1" + + +def test_resolve_workspace_owner_scope_keeps_workflow_run_for_pure_workflow() -> None: + owner = resolve_workflow_agent_workspace_owner_scope( + tenant_id="tenant-1", + app_id="app-1", + conversation_id=None, + workflow_run_id="run-1", + node_id="node-1", + workflow_agent_binding_id="workflow-binding-1", + node_execution_id="execution-1", + ) + assert owner.owner_type is AgentWorkspaceOwnerType.WORKFLOW_RUN + assert owner.owner_id == "run-1" + + +def test_conversation_scope_reuses_existing_binding_on_later_turn( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + prior_snapshot = CompositorSessionSnapshot( + layers=[ + LayerSessionSnapshot( + name="history", + lifecycle_state=LifecycleState.SUSPENDED, + runtime_state={"turn": 1}, + ) + ] + ) + workspace = AgentWorkspace( + id="workspace-conversation", + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id="conversation-1", + owner_scope_key="node-1:workflow-binding-1", + backend_workspace_ref="workspace-1-ref", + status=AgentWorkingResourceStatus.ACTIVE, + active_guard=1, + ) + prior_binding = AgentWorkspaceBinding( + id="binding-prior", + tenant_id="tenant-1", + app_id="app-1", + workspace_id=workspace.id, + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="backend-binding-prior", + status=AgentWorkingResourceStatus.ACTIVE, + session_snapshot=prior_snapshot.model_dump_json(), + ) + execution = _execution_row( + execution_id="execution-2", + workflow_run_id="run-2", + process_data={"workflow_agent_binding_id": "workflow-binding-1"}, + ) + sqlite_session.add_all([_home_snapshot(), workspace, prior_binding, execution]) + sqlite_session.commit() + client = _install_backend_client(monkeypatch) + + scope = WorkflowAgentSessionScope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-2", + node_id="node-1", + node_execution_id="execution-2", + workflow_agent_binding_id="workflow-binding-1", + agent_id="agent-1", + agent_config_snapshot_id="config-1", + conversation_id="conversation-1", + ) + stored = WorkflowAgentWorkspaceStore().load_or_create_node_execution_session(scope, home_snapshot_id="home-1") + + assert stored.workspace_id == workspace.id + assert stored.binding_id == prior_binding.id + assert stored.session_snapshot == prior_snapshot + client.create_execution_binding_sync.assert_not_called() + assert sqlite_session.scalar(select(func.count()).select_from(AgentWorkspaceBinding)) == 1 + + +def test_load_existing_scope_reads_the_generation_from_the_persisted_binding(sqlite_session: Session) -> None: + sqlite_session.add_all( + [ + _execution_row( + binding_id="binding-1", + process_data={"workflow_agent_binding_id": "workflow-binding-1"}, + ), + _workspace_row(), + _binding_row(), + ] + ) + sqlite_session.commit() + + scope = WorkflowAgentWorkspaceStore().load_existing_node_execution_scope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + ) + + assert scope is not None + assert scope.workflow_agent_binding_id == "workflow-binding-1" + assert scope.agent_id == "agent-1" + assert scope.agent_config_snapshot_id == "config-1" + + +def test_load_existing_scope_rejects_unavailable_persisted_binding(sqlite_session: Session) -> None: + sqlite_session.add( + _execution_row( + binding_id="binding-missing", + process_data={"workflow_agent_binding_id": "workflow-binding-1"}, + ) + ) + sqlite_session.commit() + + with pytest.raises(AgentWorkspaceNotFoundError, match="participant Binding is unavailable"): + WorkflowAgentWorkspaceStore().load_existing_node_execution_scope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + ) + + +@pytest.mark.parametrize("home_snapshot_id", ["home-1", None]) +def test_load_or_create_persists_binding_on_node_execution( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + home_snapshot_id: str | None, +) -> None: + execution = _execution_row(process_data={"existing": "value"}) + rows: list[object] = [execution] + if home_snapshot_id is not None: + rows.append(_home_snapshot()) + sqlite_session.add_all(rows) + sqlite_session.commit() + _install_backend_client(monkeypatch) + + stored = WorkflowAgentWorkspaceStore().load_or_create_node_execution_session( + _scope(), home_snapshot_id=home_snapshot_id + ) + + sqlite_session.expire_all() + persisted_execution = sqlite_session.get(WorkflowNodeExecutionModel, execution.id) + assert persisted_execution is not None + assert stored.workspace_id + assert stored.backend_binding_ref == "backend-binding-1" + assert persisted_execution.agent_workspace_binding_id == stored.binding_id + assert persisted_execution.process_data_dict == { + "existing": "value", + "workflow_agent_binding_id": "workflow-binding-1", + } + assert "agent_workspace_binding_id" not in persisted_execution.process_data_dict + persisted_binding = sqlite_session.get(AgentWorkspaceBinding, stored.binding_id) + assert persisted_binding is not None + assert persisted_binding.base_home_snapshot_id == home_snapshot_id + + +def test_load_existing_pointer_rejects_missing_workflow_identity(sqlite_session: Session) -> None: + execution = _execution_row(binding_id="binding-1", process_data={"existing": "value"}) + sqlite_session.add(execution) + sqlite_session.commit() + + with pytest.raises(AgentWorkspaceNotFoundError, match="caller identity is missing"): + WorkflowAgentWorkspaceStore().load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + sqlite_session.expire(execution) + assert execution.process_data_dict == {"existing": "value"} + assert sqlite_session.scalar(select(func.count()).select_from(AgentWorkspaceBinding)) == 0 + + +def test_load_existing_pointer_reuses_matching_workflow_identity(sqlite_session: Session) -> None: + original_process_data = { + "existing": "value", + "workflow_agent_binding_id": "workflow-binding-1", + } + execution = _execution_row(binding_id="binding-1", process_data=original_process_data) + sqlite_session.add_all([execution, _workspace_row(), _binding_row()]) + sqlite_session.commit() + + stored = WorkflowAgentWorkspaceStore().load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + sqlite_session.expire(execution) + assert stored.binding_id == "binding-1" + assert execution.process_data_dict == original_process_data + assert "agent_workspace_binding_id" not in execution.process_data_dict + assert sqlite_session.scalar(select(func.count()).select_from(AgentWorkspaceBinding)) == 1 + + +def test_load_existing_pointer_rejects_conflicting_workflow_identity(sqlite_session: Session) -> None: + execution = _execution_row( + binding_id="binding-1", + process_data={"workflow_agent_binding_id": "workflow-binding-other"}, + ) + sqlite_session.add(execution) + sqlite_session.commit() + + with pytest.raises(AgentWorkspaceNotFoundError, match="caller identity does not match"): + WorkflowAgentWorkspaceStore().load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + assert sqlite_session.scalar(select(func.count()).select_from(AgentWorkspaceBinding)) == 0 + + +def test_load_or_create_fails_before_binding_create_when_caller_row_is_missing( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + executed_statements: list[str], +) -> None: + sleep = MagicMock() + monkeypatch.setattr("core.workflow.nodes.agent_v2.session_store.time.sleep", sleep) + + with pytest.raises(AgentWorkspaceNotFoundError, match="Workflow node execution caller is unavailable"): + WorkflowAgentWorkspaceStore().load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + assert len(_execution_selects(executed_statements)) == 60 + assert sleep.call_count == 59 + assert sqlite_session.scalar(select(func.count()).select_from(AgentWorkspaceBinding)) == 0 + + +def test_load_existing_scope_waits_for_caller_row_to_become_visible( + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], + executed_statements: list[str], +) -> None: + sleep = MagicMock() + + def make_caller_visible(_seconds: float) -> None: + if sleep.call_count == 2: + with sqlite_session_factory() as observer: + observer.add(_execution_row()) + observer.commit() + + sleep.side_effect = make_caller_visible + monkeypatch.setattr("core.workflow.nodes.agent_v2.session_store.time.sleep", sleep) + + scope = WorkflowAgentWorkspaceStore().load_existing_node_execution_scope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + ) + + assert scope is None + assert len(_execution_selects(executed_statements)) == 3 + assert sleep.call_count == 2 + + +def test_save_snapshot_targets_binding(sqlite_session: Session) -> None: + binding = _binding_row() + sqlite_session.add_all([_workspace_row(), binding]) + sqlite_session.commit() + snapshot = CompositorSessionSnapshot(layers=[]) + + WorkflowAgentWorkspaceStore().save_active_snapshot( + scope=_scope(), + binding_id=binding.id, + snapshot=snapshot, + ) + + sqlite_session.expire(binding) + assert binding.session_snapshot == snapshot.model_dump_json() + + +def test_retire_workflow_run_only_retires_matching_tenant_and_app(sqlite_session: Session) -> None: + matching = _workspace_row() + other_tenant = _workspace_row(workspace_id="workspace-other-tenant", tenant_id="tenant-2") + other_app = _workspace_row( + workspace_id="workspace-other-app", + app_id="app-2", + owner_scope_key="node-2:workflow-binding-2", + ) + sqlite_session.add_all([matching, other_tenant, other_app]) + sqlite_session.commit() + + workspace_ids = WorkflowAgentWorkspaceStore().retire_workflow_run( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + ) + + sqlite_session.expire_all() + assert matching.status is AgentWorkingResourceStatus.RETIRED + assert other_tenant.status is AgentWorkingResourceStatus.ACTIVE + assert other_app.status is AgentWorkingResourceStatus.ACTIVE + assert workspace_ids == [matching.id] + + +def test_retire_workflow_run_transitions_active_workspace(sqlite_session: Session) -> None: + workspace = _workspace_row() + binding = _binding_row() + sqlite_session.add_all([workspace, binding]) + sqlite_session.commit() + + workspace_ids = WorkflowAgentWorkspaceStore().retire_workflow_run( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + ) + + sqlite_session.expire_all() + assert workspace.status is AgentWorkingResourceStatus.RETIRED + assert binding.status is AgentWorkingResourceStatus.RETIRED + assert workspace_ids == [workspace.id] + + +def test_retire_workflow_run_returns_existing_retired_workspace(sqlite_session: Session) -> None: + workspace = _workspace_row(status=AgentWorkingResourceStatus.RETIRED) + sqlite_session.add(workspace) + sqlite_session.commit() + + workspace_ids = WorkflowAgentWorkspaceStore().retire_workflow_run( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + ) + + sqlite_session.expire(workspace) + assert workspace.status is AgentWorkingResourceStatus.RETIRED + assert workspace_ids == [workspace.id] + + +def test_workflow_scope_creates_binding_without_existing_workspace_ref( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + execution = _execution_row(process_data={"workflow_agent_binding_id": "workflow-binding-1"}) + sqlite_session.add_all([_home_snapshot(), execution]) + sqlite_session.commit() + client = _install_backend_client(monkeypatch) + + WorkflowAgentWorkspaceStore().load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + request = client.create_execution_binding_sync.call_args.args[0] + assert request.existing_workspace_ref is None + + +def test_conversation_scope_first_turn_creates_binding_without_existing_workspace_ref( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + execution = _execution_row( + execution_id="execution-2", + workflow_run_id="run-2", + process_data={"workflow_agent_binding_id": "workflow-binding-1"}, + ) + sqlite_session.add_all([_home_snapshot(), execution]) + sqlite_session.commit() + client = _install_backend_client(monkeypatch) + + scope = WorkflowAgentSessionScope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-2", + node_id="node-1", + node_execution_id="execution-2", + workflow_agent_binding_id="workflow-binding-1", + agent_id="agent-1", + agent_config_snapshot_id="config-1", + conversation_id="conversation-1", + ) + WorkflowAgentWorkspaceStore().load_or_create_node_execution_session(scope, home_snapshot_id="home-1") + + request = client.create_execution_binding_sync.call_args.args[0] + assert request.existing_workspace_ref is None