Upstream langgenius/dify state of the touched files before 0ff8f0cd14
This commit is contained in:
upstream committed
commit
318c2c2eda
2 files changed
+936
No files matched your search
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
Reference in new issue
Block a user