diff --git a/api/extensions/application_services/human_input.py b/api/extensions/application_services/human_input.py new file mode 100644 index 00000000000000..900c12e56c9a16 --- /dev/null +++ b/api/extensions/application_services/human_input.py @@ -0,0 +1,32 @@ +"""Compose human-input debug dependencies.""" + +from functools import partial + +from sqlalchemy.orm import Session, sessionmaker + +from core.app.file_access import DatabaseFileAccessController +from extensions.application_services.workflow_variables import build_workflow_variable_service +from repositories.human_input.form_repository import HumanInputFormRepositoryImpl, HumanInputFormSubmissionRepository +from repositories.workflow.definition_repository import WorkflowDefinitionRepository +from services.human_input.debug_service import HumanInputDebugService +from services.human_input.file_gateway import HumanInputDebugFileGateway +from services.human_input_delivery_test_service import ( + DeliveryTestRegistry, + EmailDeliveryTestHandler, + HumanInputDeliveryTestService, +) +from services.human_input_service import HumanInputService + + +def build_human_input_debug_service(*, database_client: sessionmaker[Session]) -> HumanInputDebugService: + return HumanInputDebugService( + definitions=WorkflowDefinitionRepository(session_factory=database_client), + variables=build_workflow_variable_service(database_client=database_client), + files=HumanInputDebugFileGateway(sessions=database_client, access=DatabaseFileAccessController()), + submissions=HumanInputService( + session_factory=database_client, + form_repository=HumanInputFormSubmissionRepository(sessions=database_client), + ), + forms=partial(HumanInputFormRepositoryImpl, sessions=database_client), + delivery=HumanInputDeliveryTestService(DeliveryTestRegistry([EmailDeliveryTestHandler(database_client)])), + ) diff --git a/api/fields/workflow_fields.py b/api/fields/workflow_fields.py new file mode 100644 index 00000000000000..e0a8914de7e022 --- /dev/null +++ b/api/fields/workflow_fields.py @@ -0,0 +1,179 @@ +"""Shared workflow response schemas for Console, snippets and RAG pipelines.""" + +from datetime import datetime +from typing import Any, NotRequired, TypedDict + +from pydantic import AliasChoices, Field, RootModel, field_validator + +from core.helper import encrypter +from core.workflow.llm_environment_variable import LLM_ENVIRONMENT_VARIABLE_VALUE_TYPE, LLMEnvironmentVariable +from fields.base import ResponseModel +from fields.conversation_variable_fields import WorkflowConversationVariableResponse +from fields.member_fields import SimpleAccount +from graphon.variables import SecretVariable, SegmentType, VariableBase +from libs.helper import to_timestamp + +ENVIRONMENT_VARIABLE_SUPPORTED_TYPES = (SegmentType.STRING, SegmentType.NUMBER, SegmentType.SECRET) + + +class EnvironmentVariableResponseDict(TypedDict): + value_type: str + id: NotRequired[str] + name: NotRequired[str] + value: NotRequired[Any] + description: NotRequired[str | None] + + +class PipelineVariableResponse(ResponseModel): + label: str + variable: str + type: str + belong_to_node_id: str + max_length: int | None = None + required: bool + unit: str | None = None + default_value: Any = Field(default=None) + options: list[str] | None = None + placeholder: str | None = None + tooltips: str | None = None + allowed_file_types: list[str] | None = None + allowed_file_extensions: list[str] | None = Field( + default=None, validation_alias=AliasChoices("allowed_file_extensions", "allow_file_extension") + ) + allowed_file_upload_methods: list[str] | None = Field( + default=None, validation_alias=AliasChoices("allowed_file_upload_methods", "allow_file_upload_methods") + ) + + +class WorkflowEnvironmentVariableResponse(ResponseModel): + value_type: str + id: str + name: str + value: Any + description: str + + +class WorkflowResponse(ResponseModel): + id: str + graph: dict[str, Any] = Field( + validation_alias=AliasChoices("graph_dict", "graph"), + ) + features: dict[str, Any] = Field( + validation_alias=AliasChoices("features_dict", "features"), + ) + hash: str = Field(validation_alias=AliasChoices("unique_hash", "hash")) + version: str + # NULL for drafts and for versions published before numbering was introduced; those + # render as "Untitled Version" instead of `#N`. Never 0, so clients must test for null. + version_number: int | None = None + marked_name: str + marked_comment: str + created_by: SimpleAccount | None = Field( + default=None, validation_alias=AliasChoices("created_by_account", "created_by") + ) + created_at: int + updated_by: SimpleAccount | None = Field( + default=None, validation_alias=AliasChoices("updated_by_account", "updated_by") + ) + updated_at: int + tool_published: bool + environment_variables: list[WorkflowEnvironmentVariableResponse] + conversation_variables: list[WorkflowConversationVariableResponse] + rag_pipeline_variables: list[PipelineVariableResponse] + + @field_validator("created_at", "updated_at", mode="before") + @classmethod + def _normalize_timestamp(cls, value: datetime | int | None) -> int: + timestamp = to_timestamp(value) + if timestamp is None: + raise ValueError("timestamp is required") + return timestamp + + @field_validator("environment_variables", mode="before") + @classmethod + def _serialize_environment_variables(cls, value: Any) -> list[Any]: + if value is None: + return [] + + return [_serialize_environment_variable(item) for item in value] + + +class WorkflowPaginationResponse(ResponseModel): + items: list[WorkflowResponse] + page: int + limit: int + has_more: bool + + +def _serialize_environment_variable(value: Any) -> EnvironmentVariableResponseDict | Any: + match value: + case LLMEnvironmentVariable(): + return { + "id": value.id, + "name": value.name, + "value": value.value, + "value_type": LLM_ENVIRONMENT_VARIABLE_VALUE_TYPE, + "description": value.description, + } + + case SecretVariable(): + return { + "id": value.id, + "name": value.name, + "value": encrypter.full_mask_token(), + "value_type": value.value_type.value, + "description": value.description, + } + + case VariableBase(): + return { + "id": value.id, + "name": value.name, + "value": value.value, + "value_type": str(value.value_type.exposed_type()), + "description": value.description, + } + + case dict(): + value_type_str = value.get("value_type") + if not isinstance(value_type_str, str): + raise TypeError( + f"unexpected type for value_type field, value={value_type_str}, type={type(value_type_str)}" + ) + if value_type_str == LLM_ENVIRONMENT_VARIABLE_VALUE_TYPE: + return value + value_type = SegmentType(value_type_str).exposed_type() + if value_type not in ENVIRONMENT_VARIABLE_SUPPORTED_TYPES: + raise ValueError(f"Unsupported environment variable value type: {value_type}") + return value + + case _: + return value + + +class WorkflowPublishResponse(ResponseModel): + result: str + created_at: int + warning: str | None = Field( + default=None, + description="Advisory warning for variable references that can read a skipped branch. Publish still succeeds.", + ) + + +class WorkflowRestoreResponse(ResponseModel): + result: str + hash: str + updated_at: int + + +class DefaultBlockConfigItemResponse(ResponseModel): + type: str + config: dict[str, Any] + + +class DefaultBlockConfigsResponse(RootModel[list[DefaultBlockConfigItemResponse]]): + pass + + +class DefaultBlockConfigResponse(RootModel[dict[str, Any]]): + root: dict[str, Any] diff --git a/api/models/human_input_contracts.py b/api/models/human_input_contracts.py index ea9cd05fc920f1..cad6ca8436d971 100644 --- a/api/models/human_input_contracts.py +++ b/api/models/human_input_contracts.py @@ -106,6 +106,66 @@ def submitted(self) -> bool: return self.submitted_at is not None +class HumanInputFormRecordView(Protocol): + """Read-only submission state consumed by the Human Input application service.""" + + @property + def form_id(self) -> str: ... + + @property + def workflow_run_id(self) -> str | None: ... + + @property + def conversation_id(self) -> str | None: ... + + @property + def tenant_id(self) -> str: ... + + @property + def app_id(self) -> str: ... + + @property + def form_kind(self) -> HumanInputFormKind: ... + + @property + def definition(self) -> FormDefinition: ... + + @property + def created_at(self) -> datetime: ... + + @property + def expiration_time(self) -> datetime: ... + + @property + def status(self) -> HumanInputFormStatus: ... + + @property + def recipient_id(self) -> str | None: ... + + @property + def recipient_type(self) -> RecipientType | None: ... + + @property + def submitted(self) -> bool: ... + + +class HumanInputFormSubmissionStore(Protocol): + """Persistence port shared by legacy and session-bound Human Input repositories.""" + + def get_by_token(self, form_token: str) -> HumanInputFormRecordView | None: ... + + def mark_submitted( + self, + *, + form_id: str, + recipient_id: str | None, + selected_action_id: str, + form_data: Mapping[str, Any], + submission_user_id: str | None, + submission_end_user_id: str | None, + ) -> HumanInputFormRecordView: ... + + @dataclasses.dataclass(frozen=True) class PauseFormSnapshot: id: str diff --git a/api/repositories/app/agent_app_repository.py b/api/repositories/app/agent_app_repository.py index fef158cd3157ab..60247f750f259f 100644 --- a/api/repositories/app/agent_app_repository.py +++ b/api/repositories/app/agent_app_repository.py @@ -18,6 +18,11 @@ class AgentAppRepository: def __init__(self, *, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory + def app_tenant_id(self, app_id: str) -> str | None: + """Read ownership only, allowing the trusted inner API to distinguish 403 from 404.""" + with self._session_factory() as session: + return session.scalar(select(App.tenant_id).where(App.id == app_id)) + def resolve_roster_app_id(self, *, tenant_id: str, agent_id: str) -> str: """Resolve a roster Agent's configuration owner without materializing inline Apps.""" with self._session_factory() as session: diff --git a/api/repositories/credentials/query_repository.py b/api/repositories/credentials/query_repository.py index 89605c4e85d479..7d51e7e5894cf2 100644 --- a/api/repositories/credentials/query_repository.py +++ b/api/repositories/credentials/query_repository.py @@ -11,10 +11,11 @@ from models.credential_permission import CredentialPermission, CredentialType from models.enums import PermissionEnum from models.oauth import DatasourceProvider -from models.provider import ProviderCredential +from models.provider import LoadBalancingModelConfig, ProviderCredential from models.provider_ids import ModelProviderID from models.tools import BuiltinToolProvider from models.trigger import TriggerSubscription, WorkflowPluginTrigger +from repositories.tools.provider_repository import select_builtin_credentials from services.credentials.query import ( CredentialQuery, DatasourceCredentialListItem, @@ -118,6 +119,34 @@ def update_tool_credential_if_unchanged( assert isinstance(result, CursorResult) return result.rowcount == 1 + @override + def list_model_load_balancing_credential_ids( + self, *, workspace_id: str, provider: str, model: str + ) -> Sequence[str]: + # Both predefined and custom model credentials need policy validation. + # Only IDs are needed; publication must not fetch/decrypt credential forms. + names = [provider] + provider_id = ModelProviderID(provider) + if provider_id.is_langgenius(): + names.append(provider_id.provider_name if "/" in provider else str(provider_id)) + with self._session_factory() as session: + return [ + credential_id + for credential_id in session.scalars( + select(LoadBalancingModelConfig.credential_id) + .where( + LoadBalancingModelConfig.tenant_id == workspace_id, + LoadBalancingModelConfig.provider_name.in_(names), + LoadBalancingModelConfig.model_type == "llm", + LoadBalancingModelConfig.model_name == model, + LoadBalancingModelConfig.credential_id.is_not(None), + LoadBalancingModelConfig.credential_id != "", + ) + .distinct() + ) + if credential_id is not None + ] + @override def list_models(self, *, workspace_id: str, provider: str, actor_id: str | None) -> Sequence[ModelCredentialRecord]: names = [provider] @@ -185,6 +214,15 @@ def list_datasources( for row in rows ] + @override + def default_tool_credential_id(self, *, workspace_id: str, provider: str) -> str | None: + with self._session_factory() as session: + return session.scalar( + select_builtin_credentials(tenant_id=workspace_id, provider_names=[provider]) + .with_only_columns(BuiltinToolProvider.id) + .limit(1) + ) + @override def list_tools( self, diff --git a/api/repositories/workflow/node_execution_repository.py b/api/repositories/workflow/node_execution_repository.py new file mode 100644 index 00000000000000..ce1e887512c308 --- /dev/null +++ b/api/repositories/workflow/node_execution_repository.py @@ -0,0 +1,81 @@ +"""Materialized Console execution responses shared by SQL and LogStore results.""" + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from models.account import Account, TenantAccountJoin +from models.enums import CreatorUserRole +from models.model import EndUser +from models.workflow import WorkflowNodeExecutionModel + + +class WorkflowNodeExecutionRepository: + def __init__(self, session_factory: sessionmaker[Session]) -> None: + self._sessions = session_factory + + def execution_record(self, execution: WorkflowNodeExecutionModel, *, include_details: bool) -> dict[str, Any]: + """Materialize Console response data, including creators stored outside the execution backend. + + Execution/offload rows are already detached by the execution repository. + Plugin extras belong to the runtime adapter and are resolved after this read. + """ + raw = {key: value for key, value in vars(execution).items() if key in execution.__table__.columns} + if not include_details: + if "offload_data" in vars(execution): + raw["offload_data"] = [ + { + "id": offload.id, + "tenant_id": offload.tenant_id, + "app_id": offload.app_id, + "node_execution_id": offload.node_execution_id, + "type_": offload.type_.value, + "file_id": offload.file_id, + "created_at": offload.created_at, + } + for offload in execution.offload_data + ] + return raw + with self._sessions() as session: + account = ( + session.scalar( + select(Account) + .join(TenantAccountJoin, TenantAccountJoin.account_id == Account.id) + .where(Account.id == execution.created_by, TenantAccountJoin.tenant_id == execution.tenant_id) + ) + if execution.created_by_role == CreatorUserRole.ACCOUNT + else None + ) + end_user = ( + session.scalar( + select(EndUser).where( + EndUser.id == execution.created_by, + EndUser.tenant_id == execution.tenant_id, + EndUser.app_id == execution.app_id, + ) + ) + if execution.created_by_role == CreatorUserRole.END_USER + else None + ) + return { + **raw, + "inputs_dict": execution.inputs_dict, + "process_data_dict": execution.process_data_dict, + "outputs_dict": execution.outputs_dict, + "execution_metadata_dict": execution.execution_metadata_dict, + "inputs_truncated": execution.inputs_truncated, + "outputs_truncated": execution.outputs_truncated, + "process_data_truncated": execution.process_data_truncated, + "created_by_account": {"id": account.id, "name": account.name, "email": account.email} + if account is not None + else None, + "created_by_end_user": { + "id": end_user.id, + "type": end_user.type, + "is_anonymous": end_user.is_anonymous, + "session_id": end_user.session_id, + } + if end_user is not None + else None, + } diff --git a/api/services/agent/composer_service.py b/api/services/agent/composer_service.py index 2477354b394d12..2ece84a96c6ec3 100644 --- a/api/services/agent/composer_service.py +++ b/api/services/agent/composer_service.py @@ -33,6 +33,7 @@ effective_declared_outputs as _effective_declared_outputs, ) from models.workflow import Workflow +from repositories.agent.workflow_binding_repository import WorkflowAgentBindingRepository, workflow_binding_scope from services.agent.agent_soul_state import agent_soul_has_model from services.agent.composer_validator import ComposerConfigValidator from services.agent.dsl_entities import make_portable_agent_soul @@ -221,9 +222,7 @@ def save_workflow_composer( if payload.save_strategy in _PUBLISH_SAVE_STRATEGIES: cls.validate_knowledge_datasets(session=session, tenant_id=tenant_id, agent_soul=payload.agent_soul) workflow = cls._get_draft_workflow(session=session, tenant_id=tenant_id, app_id=app_id) - binding = cls._get_workflow_binding( - session=session, tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id - ) + binding = cls._lock_workflow_binding(session=session, workflow=workflow, node_id=node_id) retirement_candidates = ( {binding.agent_id} if binding is not None @@ -323,9 +322,7 @@ def copy_workflow_composer_from_roster( """Copy a callable roster Agent snapshot into a workflow-owned inline Agent.""" workflow = cls._get_draft_workflow(session=session, tenant_id=tenant_id, app_id=app_id) - binding = cls._require_binding( - cls._get_workflow_binding(session=session, tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id) - ) + binding = cls._require_binding(cls._lock_workflow_binding(session=session, workflow=workflow, node_id=node_id)) if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and idempotency_key: agent = cls._get_agent_if_present(session=session, tenant_id=tenant_id, agent_id=binding.agent_id) @@ -2326,6 +2323,20 @@ def _get_workflow_binding( stmt = stmt.with_for_update().execution_options(populate_existing=True, autoflush=False) return session.scalar(stmt.limit(1)) + @staticmethod + def _lock_workflow_binding( + *, session: Session, workflow: Workflow, node_id: str + ) -> WorkflowAgentNodeBinding | None: + """Lock every draft binding, then its Agent, before returning the requested binding. + + Debug reservation takes the same stable binding-first lock order while + materializing execution bindings. Locking the whole workflow prevents a + concurrent writer from acquiring a later binding before this transaction + reaches it. + """ + bindings = WorkflowAgentBindingRepository(session).lock_bindings(workflow_binding_scope(workflow)) + return next((binding for binding in bindings if binding.node_id == node_id), None) + @classmethod def _require_binding(cls, binding: WorkflowAgentNodeBinding | None) -> WorkflowAgentNodeBinding: if not binding: diff --git a/api/services/agent/tool_invocation_service.py b/api/services/agent/tool_invocation_service.py new file mode 100644 index 00000000000000..78c989f294974c --- /dev/null +++ b/api/services/agent/tool_invocation_service.py @@ -0,0 +1,32 @@ +"""Authorize and dispatch API-owned Agent tool invocations through explicit ports.""" + +from typing import Protocol + +from services.entities.agent_tool_inner import AgentToolInvokeRequest, AgentToolInvokeResponse +from services.errors.agent_tool_inner import AgentToolInnerServiceError + + +class ToolAppOwnership(Protocol): + def app_tenant_id(self, app_id: str) -> str | None: ... + + +class AgentToolInvocation(Protocol): + def invoke(self, request: AgentToolInvokeRequest) -> AgentToolInvokeResponse: ... + + +class AgentToolInnerService: + def __init__(self, *, apps: ToolAppOwnership, tools: AgentToolInvocation) -> None: + self._apps = apps + self._tools = tools + + def invoke(self, request: AgentToolInvokeRequest) -> AgentToolInvokeResponse: + tenant_id = self._apps.app_tenant_id(request.caller.app_id) + if tenant_id is None: + raise AgentToolInnerServiceError(error_code="app_not_found", description="App not found.", status_code=404) + if tenant_id != request.caller.tenant_id: + raise AgentToolInnerServiceError( + error_code="app_tenant_mismatch", + description="App does not belong to the caller tenant.", + status_code=403, + ) + return self._tools.invoke(request) diff --git a/api/services/credentials/query.py b/api/services/credentials/query.py index 54fb193b107e7a..4b0116df62156a 100644 --- a/api/services/credentials/query.py +++ b/api/services/credentials/query.py @@ -74,6 +74,14 @@ def list_tools( include_credential_ids: Sequence[str] = (), ) -> Sequence[ToolCredentialRecord]: ... + def default_tool_credential_id(self, *, workspace_id: str, provider: str) -> str | None: + """Return the workspace default, falling back to its oldest credential.""" + ... + def list_trigger_subscriptions( self, *, workspace_id: str, provider: str, actor_id: str | None ) -> Sequence[TriggerSubscriptionRecord]: ... + + def list_model_load_balancing_credential_ids( + self, *, workspace_id: str, provider: str, model: str + ) -> Sequence[str]: ... diff --git a/api/services/human_input/debug_service.py b/api/services/human_input/debug_service.py new file mode 100644 index 00000000000000..7dc48c9b082fc8 --- /dev/null +++ b/api/services/human_input/debug_service.py @@ -0,0 +1,234 @@ +"""Preview, submit and deliver draft human-input forms through explicit ports.""" + +from collections.abc import Mapping, Sequence +from typing import Any, Protocol + +from core.workflow.environment_variables import load_environment_variables +from core.workflow.human_input_adapter import adapt_human_input_node_data_for_graph, parse_human_input_delivery_methods +from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from core.workflow.system_variables import build_bootstrap_variables, default_system_variables +from core.workflow.variable_pool_initializer import add_variables_to_pool +from enums.human_input import HumanInputFormKind +from graphon.nodes import BuiltinNodeTypes +from graphon.runtime import VariablePool +from graphon.variable_loader import VariableLoader, load_into_variable_pool +from graphon.variables import VariableBase +from machinery.context import RequestContext +from models.human_input_contracts import FormCreateParams, HumanInputFormEntity +from models.human_input_delivery import DeliveryChannelConfig +from models.human_input_entities import HumanInputNodeData +from services.human_input.contracts import ( + DeliveryTestContext, + DeliveryTestEmailRecipient, + DeliveryTestError, + DeliveryTestResult, + DeliveryTestUnsupportedError, +) +from services.workflow.contracts import WorkflowDebugNode +from services.workflow.execution.adapters.human_input import ( + DifyHITLCallback, + render_form_content_before_submission, + resolve_default_values, +) +from services.workflow.execution.adapters.node_runtime import apply_dify_debug_email_recipient +from services.workflow_node_variables import get_human_input_form_variable_mapping + + +class DebugDefinitions(Protocol): + def debug_node(self, context: RequestContext, app_id: str, node_id: str) -> WorkflowDebugNode: ... + + +class DebugVariables(Protocol): + def loader( + self, + context: RequestContext, + app_id: str, + conversation_variables: Sequence[VariableBase], + ) -> VariableLoader: ... + def save( + self, + context: RequestContext, + app_id: str, + node_id: str, + enclosing_node_id: str | None, + outputs: Mapping[str, Any], + ) -> None: ... + + +class DebugFileInputs(Protocol): + def map_inputs( + self, + *, + tenant_id: str, + variable_mapping: Mapping[str, Sequence[str]], + user_inputs: dict[str, Any], + variable_pool: VariablePool, + ) -> None: ... + + +class FormSubmission(Protocol): + def validate_and_normalize_submission( + self, + *, + tenant_id: str, + form_definition: HumanInputNodeData, + selected_action_id: str, + form_data: Mapping[str, Any], + ) -> Mapping[str, Any]: ... + + +class DebugForms(Protocol): + def create_form(self, params: FormCreateParams) -> HumanInputFormEntity: ... + def email_recipients(self, form_id: str) -> list[tuple[str, str]]: ... + + +class DebugFormFactory(Protocol): + def __call__(self, *, tenant_id: str, app_id: str) -> DebugForms: ... + + +class DebugDelivery(Protocol): + def send_test(self, *, context: DeliveryTestContext, method: DeliveryChannelConfig) -> DeliveryTestResult: ... + + +class HumanInputDebugService: + def __init__( + self, + *, + definitions: DebugDefinitions, + variables: DebugVariables, + files: DebugFileInputs, + submissions: FormSubmission, + forms: DebugFormFactory, + delivery: DebugDelivery, + ) -> None: + self._definitions = definitions + self._variables = variables + self._files = files + self._submissions = submissions + self._forms = forms + self._delivery = delivery + + def _prepare( + self, context: RequestContext, app_id: str, node_id: str, inputs: Mapping[str, Any] + ) -> tuple[WorkflowDebugNode, HumanInputNodeData, VariablePool]: + draft = self._definitions.debug_node(context, app_id, node_id) + if draft.config["data"].get("type") != BuiltinNodeTypes.HUMAN_INPUT: + raise ValueError("Node type must be human-input.") + data = HumanInputNodeData.model_validate(adapt_human_input_node_data_for_graph(draft.config["data"])) + mapping = get_human_input_form_variable_mapping(node_id, data) + manual_inputs = dict(inputs) + # The existing loader handles nesting and manual overrides consistently + # with node execution. Repository reads finish before rendering or delivery. + loader = self._variables.loader(context, app_id, draft.conversation_variables) + environment_variables = load_environment_variables( + tenant_id=context.active_workspace_id, serialized_variables=draft.environment_variables + ) + pool = VariablePool() + add_variables_to_pool( + pool, + build_bootstrap_variables( + system_variables=default_system_variables(), environment_variables=environment_variables + ), + ) + load_into_variable_pool( + variable_loader=loader, + variable_pool=pool, + variable_mapping=mapping, + user_inputs=manual_inputs, + ) + self._files.map_inputs( + tenant_id=context.active_workspace_id, + variable_mapping=mapping, + user_inputs=manual_inputs, + variable_pool=pool, + ) + return draft, data, pool + + def preview_form( + self, context: RequestContext, app_id: str, node_id: str, inputs: dict[str, Any] + ) -> Mapping[str, Any]: + _, data, pool = self._prepare(context, app_id, node_id, inputs) + return HumanInputRequired( + form_id=node_id, + form_content=render_form_content_before_submission(data, variable_pool=pool), + inputs=data.inputs, + actions=data.user_actions, + node_id=node_id, + node_title=data.title, + resolved_default_values=resolve_default_values(data, variable_pool=pool), + ).model_dump(mode="json") + + def submit_form( + self, + context: RequestContext, + app_id: str, + node_id: str, + *, + inputs: dict[str, Any], + form_inputs: dict[str, Any], + action: str, + ) -> Mapping[str, Any]: + draft, data, pool = self._prepare(context, app_id, node_id, inputs) + normalized = self._submissions.validate_and_normalize_submission( + tenant_id=context.active_workspace_id, + form_definition=data, + selected_action_id=action, + form_data=form_inputs, + ) + selected = next((item for item in data.user_actions if item.id == action), None) + outputs: dict[str, Any] = dict(normalized) + outputs["__action_id"] = action + outputs["__action_value"] = selected.title if selected else "" + outputs["__rendered_content"] = DifyHITLCallback.render_form_content_with_outputs( + form_content=render_form_content_before_submission(data, variable_pool=pool), + outputs=outputs, + field_names=data.outputs_field_names(), + form_inputs=data.inputs, + ) + self._variables.save(context, app_id, node_id, draft.enclosing_node_id, outputs) + return outputs + + def test_delivery( + self, context: RequestContext, app_id: str, node_id: str, *, inputs: dict[str, Any], delivery_method_id: str + ) -> None: + _, data, pool = self._prepare(context, app_id, node_id, inputs) + method = next( + (item for item in parse_human_input_delivery_methods(data) if str(item.id) == delivery_method_id), None + ) + if method is None: + raise ValueError("Delivery method not found.") + method = apply_dify_debug_email_recipient(method, enabled=True, actor_id=context.account_id) + rendered = render_form_content_before_submission(data, variable_pool=pool) + forms = self._forms(tenant_id=context.active_workspace_id, app_id=app_id) + form = forms.create_form( + FormCreateParams( + workflow_execution_id=None, + node_id=node_id, + form_config=data, + rendered_content=rendered, + delivery_methods=[method], + display_in_ui=False, + resolved_default_values=resolve_default_values(data, variable_pool=pool), + form_kind=HumanInputFormKind.DELIVERY_TEST, + ) + ) + recipients = [ + DeliveryTestEmailRecipient(email=email, form_token=token) + for email, token in forms.email_recipients(form.id) + ] + delivery_context = DeliveryTestContext( + tenant_id=context.active_workspace_id, + app_id=app_id, + node_id=node_id, + node_title=data.title, + rendered_content=rendered, + template_vars={"form_id": form.id}, + recipients=recipients, + variable_pool=pool, + ) + try: + self._delivery.send_test(context=delivery_context, method=method) + except DeliveryTestUnsupportedError as exc: + raise ValueError("Delivery method does not support test send.") from exc + except DeliveryTestError as exc: + raise ValueError(str(exc)) from exc diff --git a/api/services/human_input/file_gateway.py b/api/services/human_input/file_gateway.py new file mode 100644 index 00000000000000..7c94d99a8a1c91 --- /dev/null +++ b/api/services/human_input/file_gateway.py @@ -0,0 +1,33 @@ +"""Adapt manual debug inputs to the common tenant-aware file factory.""" + +from collections.abc import Mapping, Sequence +from typing import Any + +from sqlalchemy.orm import Session, sessionmaker + +from core.app.file_access import FileAccessControllerProtocol +from factories.workflow_input import map_user_inputs_to_variable_pool +from graphon.runtime import VariablePool + + +class HumanInputDebugFileGateway: + def __init__(self, *, sessions: sessionmaker[Session], access: FileAccessControllerProtocol) -> None: + self._sessions = sessions + self._access = access + + def map_inputs( + self, + *, + tenant_id: str, + variable_mapping: Mapping[str, Sequence[str]], + user_inputs: dict[str, Any], + variable_pool: VariablePool, + ) -> None: + map_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id=tenant_id, + access_controller=self._access, + sessions=self._sessions, + ) diff --git a/api/services/human_input_delivery_test_service.py b/api/services/human_input_delivery_test_service.py index d8bdca092018bd..5e0686d66a20e5 100644 --- a/api/services/human_input_delivery_test_service.py +++ b/api/services/human_input_delivery_test_service.py @@ -1,7 +1,5 @@ from __future__ import annotations -from dataclasses import dataclass, field -from enum import StrEnum from typing import Protocol from sqlalchemy import Engine, select @@ -10,7 +8,6 @@ from configs import dify_config from extensions.ext_database import db from extensions.ext_mail import mail -from graphon.runtime import VariablePool from libs.email_template_renderer import render_email_template from models import Account, TenantAccountJoin from models.human_input_delivery import ( @@ -21,44 +18,27 @@ MemberRecipient, ) from services.feature_service import FeatureService +from services.human_input.contracts import ( + DeliveryTestContext, + DeliveryTestEmailRecipient, + DeliveryTestError, + DeliveryTestResult, + DeliveryTestStatus, + DeliveryTestUnsupportedError, +) - -class DeliveryTestStatus(StrEnum): - OK = "ok" - FAILED = "failed" - - -@dataclass(frozen=True) -class DeliveryTestEmailRecipient: - email: str - form_token: str - - -@dataclass(frozen=True) -class DeliveryTestContext: - tenant_id: str - app_id: str - node_id: str - node_title: str | None - rendered_content: str - template_vars: dict[str, str] = field(default_factory=dict) - recipients: list[DeliveryTestEmailRecipient] = field(default_factory=list) - variable_pool: VariablePool | None = None - - -@dataclass(frozen=True) -class DeliveryTestResult: - status: DeliveryTestStatus - delivered_to: list[str] = field(default_factory=list) - warnings: list[str] = field(default_factory=list) - - -class DeliveryTestError(Exception): - pass - - -class DeliveryTestUnsupportedError(DeliveryTestError): - pass +__all__ = [ + "DeliveryTestContext", + "DeliveryTestEmailRecipient", + "DeliveryTestError", + "DeliveryTestHandler", + "DeliveryTestRegistry", + "DeliveryTestResult", + "DeliveryTestStatus", + "DeliveryTestUnsupportedError", + "EmailDeliveryTestHandler", + "HumanInputDeliveryTestService", +] def _build_form_link(token: str | None) -> str | None: diff --git a/api/services/human_input_service.py b/api/services/human_input_service.py index 6feb287af3968d..6b1cd1f0265cf8 100644 --- a/api/services/human_input_service.py +++ b/api/services/human_input_service.py @@ -10,10 +10,7 @@ from configs import dify_config from core.app.file_access import DatabaseFileAccessController from core.app.layers.pause_state_persist_layer import WorkflowResumptionContext -from core.repositories.human_input_repository import ( - HumanInputFormRecord, - HumanInputFormSubmissionRepository, -) +from core.repositories.human_input_repository import HumanInputFormSubmissionRepository from core.workflow.human_input_policy import resolve_variable_select_input_options from enums.human_input import HumanInputFormKind, HumanInputFormStatus, RecipientType, ValueSourceType from factories.file_factory import build_from_mapping, build_from_mappings @@ -23,6 +20,7 @@ from graphon.runtime.runtime_state_protocol import ReadOnlyVariablePool from libs.datetime_utils import ensure_naive_utc, naive_utc_now from libs.exception import BaseHTTPException +from models.human_input_contracts import HumanInputFormRecordView, HumanInputFormSubmissionStore from models.human_input_entities import ( FileInputConfig, FileListInputConfig, @@ -48,7 +46,7 @@ class Form: - def __init__(self, record: HumanInputFormRecord): + def __init__(self, record: HumanInputFormRecordView): self._record = record def get_definition(self) -> FormDefinition: @@ -158,7 +156,7 @@ class HumanInputService: def __init__( self, session_factory: sessionmaker[Session] | Engine, - form_repository: HumanInputFormSubmissionRepository | None = None, + form_repository: HumanInputFormSubmissionStore | None = None, ): if isinstance(session_factory, Engine): session_factory = sessionmaker(bind=session_factory) diff --git a/api/services/tools/agent_configuration.py b/api/services/tools/agent_configuration.py new file mode 100644 index 00000000000000..18fbcf521406e8 --- /dev/null +++ b/api/services/tools/agent_configuration.py @@ -0,0 +1,65 @@ +"""Build agent tool configuration for responses without changing stored config.""" + +import logging +from collections.abc import Mapping +from copy import deepcopy +from typing import cast + +from pydantic import JsonValue, ValidationError + +from core.agent.entities import AgentToolEntity +from core.tools.utils.configuration import ToolParameterConfigurationManager +from models.tool_runtime_contracts import WorkflowToolQueries +from services.tools.provider_queries import ToolProviders +from services.tools.tool_manager import ToolManager + +logger = logging.getLogger(__name__) + + +def mask_agent_tool_parameters( + *, + tool_providers: ToolProviders, + workflow_queries: WorkflowToolQueries, + agent_mode: Mapping[str, JsonValue], + app_id: str, + tenant_id: str, + user_id: str, +) -> dict[str, JsonValue]: + """Return a response copy, exposing parameters only after masking succeeds.""" + result = deepcopy(dict(agent_mode)) + for tool in cast(list[JsonValue], result.get("tools") or []): + if not isinstance(tool, dict): + continue + unmasked_tool = dict(tool) + if "tool_parameters" in tool: + tool["tool_parameters"] = {} + if len(unmasked_tool) <= 3: + continue + try: + agent_tool = AgentToolEntity.model_validate(unmasked_tool) + except ValidationError: + continue + try: + tool_runtime = ToolManager.get_agent_tool_runtime( + tenant_id=tenant_id, + app_id=app_id, + agent_tool=agent_tool, + user_id=user_id, + tool_providers=tool_providers, + workflow_queries=workflow_queries, + ) + manager = ToolParameterConfigurationManager( + tenant_id=tenant_id, + tool_runtime=tool_runtime, + provider_name=agent_tool.provider_id, + provider_type=agent_tool.provider_type, + identity_id=f"AGENT.{app_id}", + ) + if agent_tool.tool_parameters: + parameters = manager.decrypt_tool_parameters(agent_tool.tool_parameters) + tool["tool_parameters"] = manager.mask_tool_parameters(parameters or {}) + else: + tool["tool_parameters"] = {} + except Exception: + logger.exception("Failed to mask agent tool parameters for tool %s", agent_tool.tool_name) + return result diff --git a/api/services/tools/agent_invocation_gateway.py b/api/services/tools/agent_invocation_gateway.py new file mode 100644 index 00000000000000..ca49d3e57ffc43 --- /dev/null +++ b/api/services/tools/agent_invocation_gateway.py @@ -0,0 +1,149 @@ +"""Adapt Agent tool declarations to the API tool runtime and its domain errors.""" + +from __future__ import annotations + +from contextlib import closing + +from core.agent.entities import AgentToolEntity +from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom +from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler +from core.tools.entities.tool_entities import ToolInvokeFrom, ToolProviderType +from core.tools.errors import ( + ToolInvokeError, + ToolNotFoundError, + ToolNotSupportedError, + ToolParameterValidationError, + ToolProviderCredentialValidationError, + ToolProviderNotFoundError, +) +from core.tools.utils.message_transformer import ToolFileMessageTransformer +from services.credentials.refresh import CredentialRefreshUnavailableError +from services.entities.agent_tool_inner import AgentToolInvokeRequest, AgentToolInvokeResponse +from services.errors.agent_tool_inner import AgentToolInnerServiceError +from services.tools.tool_engine import ToolEngine +from services.tools.tool_manager import ToolManager +from services.workflow.execution.ports import WorkflowRuntime +from services.workflow.variable_contracts import WorkflowExecutionVariables + + +class AgentToolInvocationGateway: + """Invoke one API-owned Agent tool declaration, including explicit plugin-via-core calls.""" + + def __init__(self, *, variables: WorkflowExecutionVariables, runtime: WorkflowRuntime) -> None: + self._variables = variables + self._runtime = runtime + + def invoke(self, request: AgentToolInvokeRequest) -> AgentToolInvokeResponse: + agent_tool = AgentToolEntity( + provider_type=ToolProviderType.value_of(request.tool.provider_type), + provider_id=request.tool.provider_id, + tool_name=request.tool.tool_name, + tool_parameters=dict(request.tool.runtime_parameters), + credential_id=request.tool.credential_id, + ) + try: + dify_run_context = DifyRunContext( + tenant_id=request.caller.tenant_id, + app_id=request.caller.app_id, + user_id=request.caller.user_id, + user_from=UserFrom(request.caller.user_from), + invoke_from=InvokeFrom.value_of(request.caller.invoke_from), + trace_session_id=request.caller.trace_session_id, + ) + tool_runtime = ToolManager.get_agent_tool_runtime( + tool_providers=self._runtime.tool_providers, + workflow_queries=self._runtime.tools, + workflow_runtime=self._runtime, + draft_variable_saver=self._variables.saver_factory, + tenant_id=request.caller.tenant_id, + app_id=request.caller.app_id, + agent_tool=agent_tool, + user_id=request.caller.user_id, + invoke_from=dify_run_context.invoke_from, + variable_pool=None, + allow_file_parameters=True, + use_default_for_missing_form_parameters=True, + ) + tool_runtime.runtime.dify_run_context = dify_run_context + tool_runtime.runtime.tool_invoke_from = ( + ToolInvokeFrom.WORKFLOW if request.caller.workflow_id is not None else ToolInvokeFrom.AGENT + ) + with closing( + self._runtime.tool_invoker( + tool=tool_runtime, + tool_parameters=dict(request.tool.tool_parameters), + user_id=request.caller.user_id, + workflow_tool_callback=DifyWorkflowCallbackHandler(), + workflow_call_depth=0, + conversation_id=request.caller.conversation_id, + app_id=request.caller.app_id, + ) + ) as messages: + with closing( + ToolFileMessageTransformer.transform_tool_invoke_messages( + messages=messages, + user_id=request.caller.user_id, + tenant_id=request.caller.tenant_id, + conversation_id=request.caller.conversation_id, + ) + ) as transformed: + transformed_messages = list(transformed) + except CredentialRefreshUnavailableError: + raise + except ToolProviderNotFoundError as exc: + raise AgentToolInnerServiceError( + error_code="agent_tool_declaration_not_found", + description=str(exc), + status_code=404, + ) from exc + except ToolProviderCredentialValidationError as exc: + raise AgentToolInnerServiceError( + error_code="agent_tool_credential_invalid", + description=str(exc), + status_code=422, + ) from exc + except ToolParameterValidationError as exc: + raise AgentToolInnerServiceError( + error_code="tool_parameters_invalid", + description=str(exc), + status_code=422, + ) from exc + except (ToolInvokeError, ToolNotFoundError, ToolNotSupportedError) as exc: + raise AgentToolInnerServiceError( + error_code="agent_tool_invoke_failed", + description=str(exc), + status_code=422, + ) from exc + except ValueError as exc: + raise _map_value_error(exc) from exc + except Exception as exc: + raise AgentToolInnerServiceError( + error_code="agent_tool_invoke_unexpected_error", + description=str(exc), + status_code=500, + ) from exc + + return AgentToolInvokeResponse( + messages=[message.model_dump(mode="json") for message in transformed_messages], + observation=ToolEngine.tool_response_to_str(transformed_messages), + metadata={ + "provider_type": request.tool.provider_type, + "provider_id": request.tool.provider_id, + "tool_name": request.tool.tool_name, + }, + ) + + +def _map_value_error(error: ValueError) -> AgentToolInnerServiceError: + description = str(error) + if description == "app not found": + return AgentToolInnerServiceError( + error_code="app_not_found", + description="App not found.", + status_code=404, + ) + return AgentToolInnerServiceError( + error_code="agent_tool_invoke_failed", + description=description, + status_code=422, + ) diff --git a/api/services/tools/plugin_invocation.py b/api/services/tools/plugin_invocation.py new file mode 100644 index 00000000000000..5918b91269c070 --- /dev/null +++ b/api/services/tools/plugin_invocation.py @@ -0,0 +1,68 @@ +from collections.abc import Callable, Generator +from contextlib import closing +from typing import Any + +from core.app.apps.draft_variable_saver import DraftVariableSaverFactory +from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler +from core.plugin.backwards_invocation.base import BaseBackwardsInvocation +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolProviderType +from core.tools.utils.message_transformer import ToolFileMessageTransformer +from models import Account +from services.tools.tool_manager import ToolManager +from services.workflow.execution.ports import WorkflowRuntime + + +class PluginToolBackwardsInvocation(BaseBackwardsInvocation): + """ + Backwards invocation for plugin tools. + """ + + @classmethod + def invoke_tool( + cls, + tenant_id: str, + user_id: str, + tool_type: ToolProviderType, + provider: str, + tool_name: str, + tool_parameters: dict[str, Any], + credential_id: str | None = None, + *, + workflow_runtime: WorkflowRuntime, + draft_variable_saver: Callable[[str, Account], DraftVariableSaverFactory], + ) -> Generator[ToolInvokeMessage]: + """ + invoke tool + """ + tool_runtime = ToolManager.get_tool_runtime_from_plugin( + tool_type, + tenant_id, + provider, + tool_name, + tool_parameters, + user_id=user_id, + credential_id=credential_id, + workflow_runtime=workflow_runtime, + draft_variable_saver=draft_variable_saver, + ) + + def invoke() -> Generator[ToolInvokeMessage]: + with closing( + workflow_runtime.tool_invoker( + tool=tool_runtime, + tool_parameters=tool_parameters, + user_id=user_id, + workflow_tool_callback=DifyWorkflowCallbackHandler(), + workflow_call_depth=1, + ) + ) as response: + with closing( + ToolFileMessageTransformer.transform_tool_invoke_messages( + response, + user_id=user_id, + tenant_id=tenant_id, + ) + ) as transformed: + yield from transformed + + return invoke() diff --git a/api/services/workflow/access_gateway.py b/api/services/workflow/access_gateway.py new file mode 100644 index 00000000000000..cd29a9151e35df --- /dev/null +++ b/api/services/workflow/access_gateway.py @@ -0,0 +1,31 @@ +"""Workflow collaboration visibility and asset delivery adapters.""" + +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from graphon.file import helpers as file_helpers +from machinery.context import RequestContext +from services.app.access import resolve_app_access_filter +from services.app.console_service import ConsoleAppAccess + + +class WorkflowAccessGateway: + def __init__(self, *, session_factory: sessionmaker[Session], apps: ConsoleAppAccess) -> None: + self._sessions = session_factory + self._apps = apps + + def accessible_app_ids(self, context: RequestContext, maintainers: dict[str, str | None]) -> set[str]: + if not dify_config.RBAC_ENABLED: + return set(maintainers) + access = resolve_app_access_filter(context.active_workspace_id, context.account_id) + return { + app_id + for app_id, maintainer in maintainers.items() + if access.is_app_accessible(app_id, maintainer, context.account_id) + } + + def permission_keys(self, context: RequestContext, app_id: str) -> list[str]: + return self._apps.created_permissions(context, app_id) + + def avatar_url(self, avatar: str) -> str: + return file_helpers.get_signed_file_url(avatar) diff --git a/api/services/workflow/conversion_service.py b/api/services/workflow/conversion_service.py new file mode 100644 index 00000000000000..a279939e634247 --- /dev/null +++ b/api/services/workflow/conversion_service.py @@ -0,0 +1,92 @@ +"""Convert a basic app with external preparation and one atomic persistence operation.""" + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import replace +from typing import Any, Protocol + +from core.app.app_config.entities import EasyUIBasedAppConfig +from machinery.context import RequestContext +from models.model import AppMode +from models.workflow_conversion import ConversionExtension, ConvertedWorkflow, WorkflowConversionSource +from services.entities.app_entities import AppCreationSettings, AppEvent, CreateAppParams +from services.errors.workflow_service import WorkflowConversionError + + +class WorkflowConversionApps(Protocol): + def conversion_source(self, context: RequestContext, app_id: str) -> WorkflowConversionSource: ... + def conversion_extensions(self, context: RequestContext, ids: set[str]) -> Sequence[ConversionExtension]: ... + def create_converted( + self, + context: RequestContext, + source_id: str, + params: CreateAppParams, + settings: AppCreationSettings, + workflow: ConvertedWorkflow, + ) -> AppEvent: ... + + +class AppCreatedNotifier(Protocol): + def __call__(self, *, event: AppEvent, account_id: str, backing_agent_id: str | None) -> None: ... + + +class WorkflowConversionBuilder(Protocol): + def app_config(self, source: WorkflowConversionSource) -> EasyUIBasedAppConfig: ... + def convert( + self, app_config: EasyUIBasedAppConfig, extensions: Mapping[str, ConversionExtension] + ) -> ConvertedWorkflow: ... + + +class WorkflowConversionService: + def __init__( + self, + apps: WorkflowConversionApps, + *, + converter: WorkflowConversionBuilder, + decrypt_token: Callable[[str, str], str], + notify_created: AppCreatedNotifier, + ) -> None: + self._apps = apps + self._decrypt_token = decrypt_token + self._notify_created = notify_created + self._converter = converter + + def convert(self, context: RequestContext, app_id: str, args: dict[str, Any]) -> str: + source = self._apps.conversion_source(context, app_id) + if source.mode not in {AppMode.CHAT, AppMode.COMPLETION}: + raise WorkflowConversionError(f"Current App mode: {source.mode} is not supported convert to workflow.") + if source.model_config is None: + raise WorkflowConversionError("App model config is required") + config = self._converter.app_config(source) + extension_ids = { + extension_id + for variable in config.external_data_variables + if variable.type == "api" and (extension_id := variable.config.get("api_based_extension_id")) + } + extensions = { + extension.id: replace( + extension, api_key=self._decrypt_token(context.active_workspace_id, extension.api_key) + ) + for extension in self._apps.conversion_extensions(context, extension_ids) + } + workflow = self._converter.convert(config, extensions) + params = CreateAppParams( + name=args.get("name", "Default Name") or source.name + "(workflow)", + mode="workflow" if workflow.mode == AppMode.WORKFLOW else "advanced-chat", + icon_type=args.get("icon_type", "emoji") or source.icon_type, + icon=args.get("icon", "🤖") or source.icon, + icon_background=args.get("icon_background", "#FFEAD5") or source.icon_background, + api_rpm=source.api_rpm, + api_rph=source.api_rph, + ) + settings = AppCreationSettings( + app={ + "enable_site": source.enable_site, + "enable_api": source.enable_api, + "is_demo": False, + "is_public": source.is_public, + }, + model_config=None, + ) + created = self._apps.create_converted(context, app_id, params, settings, workflow) + self._notify_created(event=created, account_id=context.account_id, backing_agent_id=None) + return created.id diff --git a/api/services/workflow/definition_gateway.py b/api/services/workflow/definition_gateway.py new file mode 100644 index 00000000000000..49c273536c44c6 --- /dev/null +++ b/api/services/workflow/definition_gateway.py @@ -0,0 +1,146 @@ +"""Adapt external definition validation and draft notifications to domain owners.""" + +from collections.abc import Mapping +from contextlib import AbstractContextManager +from typing import Any, Protocol + +from sqlalchemy.orm import Session, sessionmaker + +from events.app_event import app_draft_workflow_was_synced +from factories import variable_factory +from machinery.context import RequestContext +from models import App +from models.workflow import Workflow +from repositories.workflow.definition_repository import workflow_from_snapshot +from services.agent.workflow_contracts import AgentSkillReader, WorkflowAgentBindingStore +from services.agent.workflow_publish_service import WorkflowAgentPublishService +from services.credentials.query import CredentialQuery +from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError +from services.workflow.contracts import ( + DraftSyncCommand, + PreparedDraftSync, + PreparedEnvironmentVariables, + ValidatedWorkflowPublication, + WorkflowOwner, + WorkflowSnapshot, +) +from services.workflow.environment_variable_service import prepare_environment_variables +from services.workflow.snippet_policy import validate_snippet_graph_forbidden_nodes + + +class WorkflowDefinitionReader(Protocol): + def validation_state( + self, context: RequestContext, app_id: str, workflow_id: str | None = None + ) -> tuple[App, WorkflowSnapshot | None]: ... + def binding_reader(self) -> AbstractContextManager[WorkflowAgentBindingStore]: ... + def app(self, context: RequestContext, app_id: str) -> App: ... + + +class WorkflowDraftReader(Protocol): + def snapshot( + self, context: RequestContext, owner: WorkflowOwner, workflow_id: str | None = None + ) -> WorkflowSnapshot | None: ... + + +class WorkflowDefinitionValidator(Protocol): + def validate_features_structure(self, app_model: App, features: dict[str, Any]) -> None: ... + def validate_graph_structure(self, graph: Mapping[str, Any]) -> None: ... + def validate_publication( + self, app_model: App, draft_workflow: Workflow, *, credentials: CredentialQuery + ) -> None: ... + + +class WorkflowDefinitionGateway: + def __init__( + self, + *, + session_factory: sessionmaker[Session], + definitions: WorkflowDefinitionReader, + drafts: WorkflowDraftReader, + workflows: WorkflowDefinitionValidator, + credentials: CredentialQuery, + skills: AgentSkillReader, + ) -> None: + self._sessions = session_factory + self._definitions = definitions + self._drafts = drafts + self._workflows = workflows + self._credentials = credentials + self._skills = skills + + def prepare_sync( + self, + context: RequestContext, + owner: WorkflowOwner, + command: DraftSyncCommand, + environment: PreparedEnvironmentVariables | None, + ) -> PreparedDraftSync: + if command.environment_upserts is None and command.environment_deletions: + raise ValueError("Deleted environment variable ids require an environment variable patch.") + if owner.kind == "app": + app, draft = self._definitions.validation_state(context, owner.id) + if not command.is_collaborative or draft is None: + self._workflows.validate_features_structure(app, command.features) + self._workflows.validate_graph_structure(command.graph) + else: + draft = self._drafts.snapshot(context, owner) + if owner.kind == "snippet": + validate_snippet_graph_forbidden_nodes(command.graph) + if command.check_hash and draft is not None and draft.hash != command.unique_hash: + raise WorkflowHashNotEqualError() + if owner.kind == "snippet": + return PreparedDraftSync(draft, None) + if environment is None: + mappings = command.environment_upserts + if mappings is not None: + mappings = Workflow.normalize_environment_variable_mappings(mappings) + elif not command.is_collaborative and command.environment_variables is not None: + mappings = Workflow.normalize_environment_variable_mappings(command.environment_variables) + if mappings is not None: + environment = prepare_environment_variables( + tenant_id=context.active_workspace_id, + source=draft, + variables=[variable_factory.build_environment_variable_from_mapping(value) for value in mappings], + deleted_ids=command.environment_deletions if command.environment_upserts is not None else None, + ) + return PreparedDraftSync(draft, environment) + + def validate_publish(self, context: RequestContext, app_id: str) -> ValidatedWorkflowPublication: + app, snapshot = self._definitions.validation_state(context, app_id) + if snapshot is None: + raise ValueError("No valid workflow found.") + self._workflows.validate_publication(app, workflow_from_snapshot(snapshot), credentials=self._credentials) + with self._definitions.binding_reader() as bindings: + agents = WorkflowAgentPublishService(repository=bindings).publication_state(draft_workflow=snapshot) + WorkflowAgentPublishService.validate_publication_state(agents, skills=self._skills) + return ValidatedWorkflowPublication(snapshot, agents) + + def validate_restore(self, context: RequestContext, owner: WorkflowOwner, workflow_id: str) -> WorkflowSnapshot: + if owner.kind == "app": + app, snapshot = self._definitions.validation_state(context, owner.id, workflow_id) + else: + app = None + snapshot = self._drafts.snapshot(context, owner, workflow_id) + if snapshot is None: + raise WorkflowNotFoundError("Workflow not found.") + if snapshot.version == Workflow.VERSION_DRAFT: + raise IsDraftWorkflowError("source workflow must be published") + if app is not None: + workflow = workflow_from_snapshot(snapshot) + self._workflows.validate_features_structure(app, workflow.normalized_features_dict) + self._workflows.validate_graph_structure(workflow.graph_dict) + elif owner.kind == "snippet": + validate_snippet_graph_forbidden_nodes(snapshot.graph_dict) + return snapshot + + def validate_features(self, context: RequestContext, app_id: str, features: dict[str, Any]) -> None: + app = self._definitions.app(context, app_id) + self._workflows.validate_features_structure(app, features) + + def draft_synced(self, context: RequestContext, owner: WorkflowOwner, workflow: WorkflowSnapshot) -> None: + if owner.kind != "app": + return + app = self._definitions.app(context, owner.id) + app_draft_workflow_was_synced.send( + app, synced_draft_workflow=workflow_from_snapshot(workflow), session_factory=self._sessions + ) diff --git a/api/services/workflow/node_execution_extras.py b/api/services/workflow/node_execution_extras.py new file mode 100644 index 00000000000000..03a580981caabc --- /dev/null +++ b/api/services/workflow/node_execution_extras.py @@ -0,0 +1,46 @@ +"""Resolve external provider presentation data after execution reads close.""" + +from typing import Any + +from core.tools.entities.tool_entities import ToolProviderType +from core.trigger.constants import TRIGGER_PLUGIN_NODE_TYPE +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey +from models.workflow import WorkflowNodeExecutionModel +from services.tools.provider_queries import ToolProviderIcons + + +def node_execution_extras( + execution: WorkflowNodeExecutionModel, *, tool_providers: ToolProviderIcons +) -> dict[str, Any]: + from core.trigger.trigger_manager import TriggerManager + from services.tools.tool_manager import ToolManager + + extras: dict[str, Any] = {} + execution_metadata = execution.execution_metadata_dict + if execution_metadata: + if execution.node_type == BuiltinNodeTypes.TOOL and "tool_info" in execution_metadata: + tool_info: dict[str, Any] = execution_metadata["tool_info"] + provider_type = ToolProviderType(tool_info["provider_type"]) + if provider_type == ToolProviderType.WORKFLOW: + extras["workflow_tool"] = True + extras["icon"] = ToolManager.get_tool_icon( + tenant_id=execution.tenant_id, + provider_type=provider_type, + provider_id=tool_info["provider_id"], + tool_providers=tool_providers, + ) + elif execution.node_type == BuiltinNodeTypes.DATASOURCE and "datasource_info" in execution_metadata: + datasource_info = execution_metadata["datasource_info"] + extras["icon"] = datasource_info.get("icon") + elif ( + execution.node_type == TRIGGER_PLUGIN_NODE_TYPE + and WorkflowNodeExecutionMetadataKey.TRIGGER_INFO in execution_metadata + ): + trigger_info = execution_metadata[WorkflowNodeExecutionMetadataKey.TRIGGER_INFO] or {} + provider_id = trigger_info.get("provider_id") + if provider_id: + extras["icon"] = TriggerManager.get_trigger_plugin_icon( + tenant_id=execution.tenant_id, + provider_id=provider_id, + ) + return extras diff --git a/api/tests/e2e_tests/containers/core/repositories/test_human_input_form_repository_impl.py b/api/tests/e2e_tests/containers/core/repositories/test_human_input_form_repository_impl.py index 754e2bcea45b0a..b01703303dda9b 100644 --- a/api/tests/e2e_tests/containers/core/repositories/test_human_input_form_repository_impl.py +++ b/api/tests/e2e_tests/containers/core/repositories/test_human_input_form_repository_impl.py @@ -5,9 +5,8 @@ from uuid import uuid4 from sqlalchemy import Engine, select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker -from core.repositories.human_input_repository import FormCreateParams, HumanInputFormRepositoryImpl from enums.human_input import RecipientType from models.account import ( Account, @@ -23,16 +22,18 @@ HumanInputForm, HumanInputFormRecipient, ) +from models.human_input_contracts import FormCreateParams from models.human_input_delivery import ( DeliveryChannelConfig, EmailDeliveryConfig, EmailDeliveryMethod, EmailRecipients, ExternalRecipient, + InteractiveSurfaceDeliveryMethod, MemberRecipient, ) -from models.human_input_delivery import InteractiveSurfaceDeliveryMethod as WebAppDeliveryMethod from models.human_input_entities import FormDefinition, HumanInputNodeData, UserActionConfig +from repositories.human_input.form_repository import HumanInputFormRepositoryImpl def _create_tenant_with_members(session: Session, member_emails: list[str]) -> tuple[Tenant, list[Account]]: @@ -103,7 +104,9 @@ def test_create_form_with_whole_workspace_recipients(self, db_session_with_conta member_emails=["member1@example.com", "member2@example.com"], ) - repository = HumanInputFormRepositoryImpl(tenant_id=tenant.id, app_id=str(uuid4())) + repository = HumanInputFormRepositoryImpl( + tenant_id=tenant.id, app_id=str(uuid4()), sessions=sessionmaker(bind=engine) + ) params = _build_form_params( delivery_methods=[_build_email_delivery(whole_workspace=True, recipients=[])], ) @@ -132,7 +135,9 @@ def test_create_form_with_specific_members_and_external(self, db_session_with_co member_emails=["primary@example.com", "secondary@example.com"], ) - repository = HumanInputFormRepositoryImpl(tenant_id=tenant.id, app_id=str(uuid4())) + repository = HumanInputFormRepositoryImpl( + tenant_id=tenant.id, app_id=str(uuid4()), sessions=sessionmaker(bind=engine) + ) params = _build_form_params( delivery_methods=[ _build_email_delivery( @@ -176,7 +181,9 @@ def test_create_form_persists_default_values(self, db_session_with_containers: S member_emails=["prefill@example.com"], ) - repository = HumanInputFormRepositoryImpl(tenant_id=tenant.id, app_id=str(uuid4())) + repository = HumanInputFormRepositoryImpl( + tenant_id=tenant.id, app_id=str(uuid4()), sessions=sessionmaker(bind=engine) + ) resolved_values = {"greeting": "Hello!"} params = FormCreateParams( workflow_execution_id=str(uuid4()), @@ -212,7 +219,9 @@ def test_create_form_persists_display_in_ui(self, db_session_with_containers: Se member_emails=["ui@example.com"], ) - repository = HumanInputFormRepositoryImpl(tenant_id=tenant.id, app_id=str(uuid4())) + repository = HumanInputFormRepositoryImpl( + tenant_id=tenant.id, app_id=str(uuid4()), sessions=sessionmaker(bind=engine) + ) params = FormCreateParams( workflow_execution_id=str(uuid4()), node_id="human-input-node", @@ -221,10 +230,10 @@ def test_create_form_persists_display_in_ui(self, db_session_with_containers: Se form_content="

Approve?

", inputs=[], user_actions=[UserActionConfig(id="approve", title="Approve")], - delivery_methods=[WebAppDeliveryMethod()], + delivery_methods=[InteractiveSurfaceDeliveryMethod()], ), rendered_content="

Approve?

", - delivery_methods=[WebAppDeliveryMethod()], + delivery_methods=[InteractiveSurfaceDeliveryMethod()], display_in_ui=True, resolved_default_values={}, ) diff --git a/api/tests/helpers/unit/credential_query_fakes.py b/api/tests/helpers/unit/credential_query_fakes.py index 44a7e7c6889ec4..d85156097b1e6b 100644 --- a/api/tests/helpers/unit/credential_query_fakes.py +++ b/api/tests/helpers/unit/credential_query_fakes.py @@ -41,6 +41,11 @@ def list_datasources( def list_models(self, *, workspace_id: str, provider: str, actor_id: str | None) -> Sequence[ModelCredentialRecord]: raise AssertionError(f"unexpected model credential lookup: {workspace_id=} {provider=} {actor_id=}") + def list_model_load_balancing_credential_ids( + self, *, workspace_id: str, provider: str, model: str + ) -> Sequence[str]: + raise AssertionError(f"unexpected model load-balancing credential lookup: {workspace_id=} {provider=} {model=}") + def list_tools( self, *, @@ -53,6 +58,9 @@ def list_tools( f"unexpected tool credential lookup: {workspace_id=} {provider=} {actor_id=} {include_credential_ids=}" ) + def default_tool_credential_id(self, *, workspace_id: str, provider: str) -> str | None: + raise AssertionError(f"unexpected default tool credential lookup: {workspace_id=} {provider=}") + def list_trigger_subscriptions( self, *, workspace_id: str, provider: str, actor_id: str | None ) -> Sequence[TriggerSubscriptionRecord]: diff --git a/api/tests/test_containers_integration_tests/repositories/agent/test_workflow_binding_locking.py b/api/tests/test_containers_integration_tests/repositories/agent/test_workflow_binding_locking.py new file mode 100644 index 00000000000000..d2dbd8fad415ce --- /dev/null +++ b/api/tests/test_containers_integration_tests/repositories/agent/test_workflow_binding_locking.py @@ -0,0 +1,207 @@ +"""PostgreSQL row-lock regression for concurrent Composer saves and debug reservations.""" + +import json +import time +from concurrent.futures import ThreadPoolExecutor +from queue import Queue +from threading import Event +from uuid import uuid4 + +import pytest +from sqlalchemy import Connection, Engine, event, select, text +from sqlalchemy.orm import Session, sessionmaker + +from enums.agent import WorkflowAgentBindingType +from machinery.context import RequestContext +from models.account import Account, Tenant +from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource, WorkflowAgentNodeBinding +from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig +from models.model import App +from models.workflow import Workflow, WorkflowRun +from repositories.workflow.debug_reservation_repository import WorkflowDebugReservationRepository +from services.agent.composer_service import AgentComposerService +from services.entities.agent_entities import ( + ComposerSavePayload, + ComposerSaveStrategy, + ComposerSoulLockPayload, + ComposerVariant, +) +from services.workflow.contracts import DebugReservationCursor, WorkflowSnapshot +from tests.unit_tests.workflow_execution import debug_lease + + +@pytest.mark.parametrize( + "strategy", [ComposerSaveStrategy.SAVE_AS_NEW_VERSION, ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION] +) +def test_debug_reservation_and_composer_use_the_same_lock_order( + db_session_with_containers: Session, strategy: ComposerSaveStrategy +) -> None: + session = db_session_with_containers + engine = session.get_bind() + assert isinstance(engine, Engine) + assert engine.dialect.name == "postgresql", "This regression requires real PostgreSQL row locks" + tenant = Tenant(name="Concurrent workflow") + account = Account(name="Editor", email=f"{uuid4()}@example.com") + session.add_all([tenant, account]) + session.flush() + app = App( + tenant_id=tenant.id, + name="Workflow", + mode="workflow", + created_by=account.id, + enable_site=True, + enable_api=True, + ) + session.add(app) + session.flush() + workflow = Workflow.new( + tenant_id=tenant.id, + app_id=app.id, + type="workflow", + version="draft", + features="{}", + graph=json.dumps({"nodes": [{"id": "agent", "data": {"type": "agent", "version": "2"}}], "edges": []}), + created_by=account.id, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + session.add(workflow) + session.flush() + agent = Agent( + tenant_id=tenant.id, + name="Inline", + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + app_id=app.id, + workflow_id=workflow.id, + workflow_node_id="agent", + created_by=account.id, + ) + session.add(agent) + session.flush() + soul = AgentSoulConfig.model_validate( + {"model": {"plugin_id": "langgenius/openai/openai", "model_provider": "openai", "model": "gpt-4o"}} + ) + snapshot = AgentConfigSnapshot( + tenant_id=tenant.id, + agent_id=agent.id, + version=1, + config_snapshot=soul, + created_by=account.id, + ) + session.add(snapshot) + session.flush() + agent.active_config_snapshot_id = snapshot.id + session.add( + WorkflowAgentNodeBinding( + tenant_id=tenant.id, + app_id=app.id, + workflow_id=workflow.id, + workflow_version="draft", + node_id="agent", + node_job_config=WorkflowNodeJobConfig(), + agent_id=agent.id, + current_snapshot_id=snapshot.id, + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + ) + ) + session.commit() + context = RequestContext(str(uuid4()), None, account.id, tenant.id) + app_id, agent_id = app.id, agent.id + session.rollback() # Release fixture reads before concurrent writers begin. + sessions = sessionmaker(bind=engine, expire_on_commit=False) + agent_updated, release_composer = Event(), Event() + debug_pid: Queue[int] = Queue() + + def set_up_connection(session: Session, _transaction: object, connection: Connection) -> None: + connection.info["lock_test_worker"] = session.info.get("lock_test_worker") + connection.execute(text("SET LOCAL lock_timeout = '5s'")) + if session.info.get("lock_test_worker") == "debug": + pid = connection.scalar(text("SELECT pg_backend_pid()")) + assert isinstance(pid, int) + debug_pid.put(pid) + + def hold_agent_update( + connection: Connection, + _cursor: object, + statement: str, + _parameters: object, + _context: object, + _executemany: bool, + ) -> None: + if connection.info.get("lock_test_worker") == "composer" and statement.lstrip().lower().startswith( + "update agents" + ): + agent_updated.set() + assert release_composer.wait(timeout=10), "test did not release Composer" + + event.listen(sessions, "after_begin", set_up_connection) + event.listen(engine, "after_cursor_execute", hold_agent_update) + + def save() -> None: + with sessions(info={"lock_test_worker": "composer"}) as editor: + AgentComposerService.save_workflow_composer( + session=editor, + tenant_id=context.active_workspace_id, + app_id=app_id, + node_id="agent", + account_id=context.account_id, + payload=ComposerSavePayload( + variant=ComposerVariant.WORKFLOW, + save_strategy=strategy, + soul_lock=ComposerSoulLockPayload(locked=False), + agent_soul=soul, + ), + ) + + def reserve() -> tuple[App, WorkflowSnapshot]: + debug_sessions = sessionmaker(bind=engine, expire_on_commit=False, info={"lock_test_worker": "debug"}) + event.listen(debug_sessions, "after_begin", set_up_connection) + return WorkflowDebugReservationRepository(debug_sessions).reserve_trigger_debug(context, app_id) + + try: + with ThreadPoolExecutor(max_workers=2) as workers: + saving = workers.submit(save) + try: + if not agent_updated.wait(timeout=5): + if saving.done(): + saving.result() + pytest.fail("Composer completed without updating its Agent") + pytest.fail("Composer did not reach its Agent update") + debugging = workers.submit(reserve) + pid = debug_pid.get(timeout=5) + deadline = time.monotonic() + 5 + # Observe an actual database wait, rather than assuming a sleep + # produced the desired ordering. The old code waits on Agent + # while holding binding; the fixed code waits on binding first. + while time.monotonic() < deadline: + with engine.connect() as observer: + blocked = observer.scalar(text("SELECT cardinality(pg_blocking_pids(:pid)) > 0"), {"pid": pid}) + if blocked: + break + time.sleep(0.01) + else: + pytest.fail("Debug reservation did not contend with the Composer transaction") + finally: + release_composer.set() + saving.result(timeout=10) + _, reserved = debugging.result(timeout=10) + with sessions() as reader: + current_snapshot = reader.scalar(select(Agent.active_config_snapshot_id).where(Agent.id == agent_id)) + assert reserved.graph_dict["_agent_bindings"]["bindings"]["agent"]["current_snapshot_id"] == current_snapshot + assert reserved.execution_id is not None + with sessions() as reader: + lease = debug_lease(reader, reserved.execution_id) + assert lease is not None + run = reader.get(WorkflowRun, reserved.execution_id) + assert run is not None + assert lease.expires_at is not None + cursor = DebugReservationCursor(run.id, lease.expires_at) + reservations = WorkflowDebugReservationRepository(sessions) + assert reservations.pending_batch(lease.expires_at, limit=1) == [cursor] + assert reservations.pending_batch(lease.expires_at, limit=1, after=cursor) == [] + finally: + release_composer.set() + event.remove(engine, "after_cursor_execute", hold_agent_update) + event.remove(sessions, "after_begin", set_up_connection) diff --git a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py index 9c18893fe62b46..39d3e42582dcf3 100644 --- a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py +++ b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py @@ -560,7 +560,10 @@ def test_workflow_tool_detail_hint_in_live_and_saved_traces( execution_metadata=json.dumps({"tool_info": {"provider_type": provider_type, "provider_id": "provider-1"}}), ) assert response is not None + from services.workflow.node_execution_extras import node_execution_extras + assert response.data.extras == saved.extras + assert response.data.extras == node_execution_extras(saved, tool_providers=tool_providers) assert response.data.extras.get("workflow_tool", False) == (provider_type == "workflow") @staticmethod diff --git a/api/tests/unit_tests/core/repositories/test_human_input_repository.py b/api/tests/unit_tests/core/repositories/test_human_input_repository.py index 15f4c9a51d08e6..9571d68c251531 100644 --- a/api/tests/unit_tests/core/repositories/test_human_input_repository.py +++ b/api/tests/unit_tests/core/repositories/test_human_input_repository.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from collections.abc import Iterator, Sequence +from collections.abc import Callable, Iterator, Sequence from datetime import datetime, timedelta from typing import Any @@ -11,17 +11,6 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session, sessionmaker -from core.repositories.human_input_repository import ( - FormCreateParams, - FormNotFoundError, - HumanInputFormRecord, - HumanInputFormRepositoryImpl, - HumanInputFormSubmissionRepository, - _HumanInputFormEntityImpl, - _HumanInputFormRecipientEntityImpl, - _InvalidTimeoutStatusError, - _WorkspaceMemberInfo, -) from enums.human_input import ( DeliveryMethodType, HumanInputFormKind, @@ -45,6 +34,7 @@ InstantMessageRecipientPayload, StandaloneWebAppRecipientPayload, ) +from models.human_input_contracts import FormCreateParams, FormNotFoundError, HumanInputFormRecord from models.human_input_delivery import ( EmailDeliveryConfig, EmailDeliveryMethod, @@ -55,32 +45,29 @@ InstantMessageDeliveryMethod, InstantMessageRecipients, InstantMessageUserRecipient, + InteractiveSurfaceDeliveryMethod, MemberRecipient, ) -from models.human_input_delivery import InteractiveSurfaceDeliveryMethod as WebAppDeliveryMethod from models.human_input_entities import FormDefinition, HumanInputNodeData, UserActionConfig from repositories.human_input.form_repository import ( - FormNotFoundError as _ExtractedFormNotFoundError, -) -from repositories.human_input.form_repository import ( - HumanInputFormRecord as _ExtractedHumanInputFormRecord, -) -from repositories.human_input.form_repository import ( - HumanInputFormRepositoryImpl as _ExtractedHumanInputFormRepositoryImpl, -) -from repositories.human_input.form_repository import ( - HumanInputFormSubmissionRepository as _ExtractedHumanInputFormSubmissionRepository, -) -from repositories.human_input.form_repository import ( - _WorkspaceMemberInfo as _ExtractedWorkspaceMemberInfo, + HumanInputFormRepositoryImpl, + HumanInputFormSubmissionRepository, + _HumanInputFormEntityImpl, + _HumanInputFormRecipientEntityImpl, + _InvalidTimeoutStatusError, + _WorkspaceMemberInfo, ) from repositories.human_input_errors import FormAlreadyHandledError """Unit tests for HumanInputFormRepositoryImpl private helpers.""" -def _build_repository() -> _ExtractedHumanInputFormRepositoryImpl: - return _ExtractedHumanInputFormRepositoryImpl(tenant_id="tenant-id") +def _sessions(session: Session) -> Callable[[], Session]: + return sessionmaker(bind=session.get_bind(), expire_on_commit=False) + + +def _build_repository(session: Session) -> HumanInputFormRepositoryImpl: + return HumanInputFormRepositoryImpl(tenant_id="tenant-id", sessions=_sessions(session)) def _add_workspace_member( @@ -171,11 +158,11 @@ def test_build_email_recipients_with_member_and_external( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, ) -> None: - repo = _build_repository() + repo = _build_repository(sqlite_session) _add_workspace_member(sqlite_session, user_id="member-1", email="member@example.com") query_workspace_members_by_ids = repo._query_workspace_members_by_ids - def query_members(session: Session, restrict_to_user_ids: list[str]) -> list[_ExtractedWorkspaceMemberInfo]: + def query_members(session: Session, restrict_to_user_ids: list[str]) -> list[_WorkspaceMemberInfo]: assert session is sqlite_session return query_workspace_members_by_ids( session=session, @@ -209,7 +196,7 @@ def query_members(session: Session, restrict_to_user_ids: list[str]) -> list[_Ex assert external_payload.email == "external@example.com" def test_build_email_recipients_skips_unknown_members(self, sqlite_session: Session) -> None: - repo = _build_repository() + repo = _build_repository(sqlite_session) recipients = repo._build_email_recipients( session=sqlite_session, @@ -228,7 +215,7 @@ def test_build_email_recipients_skips_unknown_members(self, sqlite_session: Sess assert recipients[0].recipient_type == RecipientType.EMAIL_EXTERNAL def test_build_email_recipients_whole_workspace_uses_all_members(self, sqlite_session: Session) -> None: - repo = _build_repository() + repo = _build_repository(sqlite_session) _add_workspace_member(sqlite_session, user_id="member-1", email="member1@example.com") _add_workspace_member(sqlite_session, user_id="member-2", email="member2@example.com") @@ -247,7 +234,7 @@ def test_build_email_recipients_whole_workspace_uses_all_members(self, sqlite_se assert emails == {"member1@example.com", "member2@example.com"} def test_build_email_recipients_dedupes_external_by_email(self, sqlite_session: Session) -> None: - repo = _build_repository() + repo = _build_repository(sqlite_session) recipients = repo._build_email_recipients( session=sqlite_session, @@ -265,7 +252,7 @@ def test_build_email_recipients_dedupes_external_by_email(self, sqlite_session: assert len(recipients) == 1 def test_build_email_recipients_prefers_member_over_external_by_email(self, sqlite_session: Session) -> None: - repo = _build_repository() + repo = _build_repository(sqlite_session) _add_workspace_member(sqlite_session, user_id="member-1", email="shared@example.com") recipients = repo._build_email_recipients( @@ -288,7 +275,7 @@ def test_delivery_method_to_model_includes_external_recipients_with_whole_worksp self, sqlite_session: Session, ) -> None: - repo = _build_repository() + repo = _build_repository(sqlite_session) _add_workspace_member(sqlite_session, user_id="member-1", email="member1@example.com") _add_workspace_member(sqlite_session, user_id="member-2", email="member2@example.com") @@ -325,11 +312,14 @@ def test_mark_timeout_rejects_form_outside_execution_scope( form = _make_form(expiration_time=naive_utc_now() - timedelta(minutes=1)) setattr(form, mismatched_scope, "other-value") _human_input_form_repository_impl_persist_form(sqlite_session, form) - repo = _ExtractedHumanInputFormRepositoryImpl( - tenant_id="tenant-id", app_id="app-id", workflow_execution_id="run-1" + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant-id", + app_id="app-id", + workflow_execution_id="run-1", + sessions=_sessions(sqlite_session), ) - with pytest.raises(_ExtractedFormNotFoundError): + with pytest.raises(FormNotFoundError): repo.mark_timeout("node-1", form_id="form-1") sqlite_session.refresh(form) @@ -344,8 +334,11 @@ def test_mark_timeout_preserves_waiting_form_that_is_not_eligible( submitted_at=naive_utc_now() if already_submitted else None, ) _human_input_form_repository_impl_persist_form(sqlite_session, form) - repo = _ExtractedHumanInputFormRepositoryImpl( - tenant_id="tenant-id", app_id="app-id", workflow_execution_id="run-1" + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant-id", + app_id="app-id", + workflow_execution_id="run-1", + sessions=_sessions(sqlite_session), ) entity = repo.mark_timeout("node-1", form_id="form-1") @@ -361,8 +354,10 @@ def test_get_form_returns_entity_and_recipients(self, sqlite_session: Session): other_tenant_form = _make_form(form_id="form-2", tenant_id="other-tenant") sqlite_session.add(other_tenant_form) _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) - repo = _ExtractedHumanInputFormRepositoryImpl( - tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant-id", + workflow_execution_id=form.workflow_run_id, + sessions=_sessions(sqlite_session), ) entity = repo.get_form(form.node_id) @@ -372,16 +367,22 @@ def test_get_form_returns_entity_and_recipients(self, sqlite_session: Session): assert len(entity.recipients) == 1 assert entity.recipients[0].token == "token-123" - def test_get_form_returns_none_when_missing(self): - repo = _ExtractedHumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id="run-1") + def test_get_form_returns_none_when_missing(self, sqlite_session: Session) -> None: + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant-id", + workflow_execution_id="run-1", + sessions=_sessions(sqlite_session), + ) assert repo.get_form("node-1") is None def test_get_form_returns_unsubmitted_state(self, sqlite_session: Session): form = _make_form() _human_input_form_repository_impl_persist_form(sqlite_session, form) - repo = _ExtractedHumanInputFormRepositoryImpl( - tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant-id", + workflow_execution_id=form.workflow_run_id, + sessions=_sessions(sqlite_session), ) entity = repo.get_form(form.node_id) @@ -398,8 +399,10 @@ def test_get_form_returns_submission_when_completed(self, sqlite_session: Sessio submitted_at=naive_utc_now(), ) _human_input_form_repository_impl_persist_form(sqlite_session, form) - repo = _ExtractedHumanInputFormRepositoryImpl( - tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant-id", + workflow_execution_id=form.workflow_run_id, + sessions=_sessions(sqlite_session), ) entity = repo.get_form(form.node_id) @@ -414,7 +417,7 @@ def test_get_by_token_returns_record(self, sqlite_session: Session): form = _make_form(tenant_id="tenant-1") recipient = _make_recipient(form.id) _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) - repo = _ExtractedHumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(sqlite_session)) record = repo.get_by_token("token-123") @@ -427,7 +430,7 @@ def test_get_by_form_id_and_recipient_type_uses_recipient(self, sqlite_session: form = _make_form(tenant_id="tenant-1") recipient = _make_recipient(form.id) _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) - repo = _ExtractedHumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(sqlite_session)) record = repo.get_by_form_id_and_recipient_type( form_id=form.id, @@ -449,9 +452,9 @@ def test_mark_submitted_updates_fields( form = _make_form(tenant_id="tenant-1", expiration_time=fixed_now) recipient = _make_recipient(form.id) _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) - repo = _ExtractedHumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(sqlite_session)) - record: _ExtractedHumanInputFormRecord = repo.mark_submitted( + record: HumanInputFormRecord = repo.mark_submitted( form_id=form.id, recipient_id=recipient.id, selected_action_id="approve", @@ -480,7 +483,7 @@ def test_mark_submitted_rolls_back_on_database_failure( form = _make_form(tenant_id="tenant-1") recipient = _make_recipient(form.id) _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) - repo = _ExtractedHumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(sqlite_session)) def fail_form_update(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: if statement.lstrip().upper().startswith("UPDATE HUMAN_INPUT_FORMS"): @@ -577,6 +580,23 @@ def _persist_recipient( class TestHumanInputRepositoryCases: + @staticmethod + @pytest.fixture + def repository_session(sqlite_engine: Engine) -> Iterator[Session]: + """Bind repository-owned sessions to an isolated SQLite database.""" + + tables = [ + HumanInputForm.__table__, + HumanInputDelivery.__table__, + HumanInputFormRecipient.__table__, + _human_input_repository_Account.__table__, + _human_input_repository_TenantAccountJoin.__table__, + ] + TypeBase.metadata.create_all(sqlite_engine, tables=tables) + repository_session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + with repository_session_factory() as session: + yield session + @staticmethod def test_recipient_entity_token_raises_when_missing() -> None: recipient = HumanInputFormRecipient( @@ -599,33 +619,6 @@ def test_recipient_entity_id_and_token_success(repository_session: Session) -> N assert entity.id == "r1" assert entity.token == "tok" - @staticmethod - def test_form_entity_submission_token_prefers_console_then_webapp_then_none(repository_session: Session) -> None: - form = _human_input_repository_persist_form(repository_session, form_id="f1") - console = _persist_recipient( - repository_session, - form_id=form.id, - recipient_id="c1", - recipient_type=RecipientType.CONSOLE, - access_token="ctok", - ) - webapp = _persist_recipient( - repository_session, - form_id=form.id, - recipient_id="w1", - recipient_type=RecipientType.STANDALONE_WEB_APP, - access_token="wtok", - ) - - entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[webapp, console]) - assert entity.submission_token == "ctok" - - entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[webapp]) - assert entity.submission_token == "wtok" - - entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[]) - assert entity.submission_token is None - @staticmethod def test_form_entity_submitted_data_parsed(repository_session: Session) -> None: form = _human_input_repository_persist_form(repository_session, form_id="f1") @@ -647,7 +640,8 @@ def test_form_record_from_models_injects_expiration_time_when_missing(repository form.expiration_time = expiration form.submitted_data = '{"k": "v"}' repository_session.commit() - record = HumanInputFormRecord.from_models(form, None) + record = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)).get_by_form_id(form.id) + assert record is not None assert record.definition.expiration_time == expiration assert record.submitted_data == {"k": "v"} assert record.submitted is False @@ -730,9 +724,9 @@ def test_delivery_method_to_model_webapp_creates_delivery_and_recipient( repository_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: "del-1") + monkeypatch.setattr("repositories.human_input.form_repository.uuidv7", lambda: "del-1") result = repo._delivery_method_to_model( - session=repository_session, form_id="form-1", delivery_method=WebAppDeliveryMethod() + session=repository_session, form_id="form-1", delivery_method=InteractiveSurfaceDeliveryMethod() ) assert result.delivery.id == "del-1" assert result.delivery.form_id == "form-1" @@ -744,7 +738,7 @@ def test_delivery_method_to_model_email_uses_build_email_recipients( repository_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: "del-1") + monkeypatch.setattr("repositories.human_input.form_repository.uuidv7", lambda: "del-1") called: dict[str, Any] = {} def fake_build(*, session: Session, form_id: str, delivery_id: str, recipients_config: Any) -> list[Any]: @@ -821,13 +815,21 @@ def fake_query(*, session: Any, restrict_to_user_ids: Sequence[str]) -> list[_Wo @staticmethod def test_get_form_returns_entity_and_none_when_missing(repository_session: Session) -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant", workflow_execution_id="run") + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant", + workflow_execution_id="run", + sessions=_sessions(repository_session), + ) assert repo.get_form("node") is None form = _human_input_repository_persist_form(repository_session, form_id="f1") _human_input_repository_persist_form(repository_session, form_id="other-tenant-form", tenant_id="other-tenant") _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") - repo = HumanInputFormRepositoryImpl(tenant_id="tenant", workflow_execution_id="run") + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant", + workflow_execution_id="run", + sessions=_sessions(repository_session), + ) entity = repo.get_form("node") assert entity is not None assert entity.id == "f1" @@ -839,10 +841,10 @@ def test_create_form_adds_console_and_backstage_recipients( repository_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: fixed_now = datetime(2024, 1, 1, 0, 0, 0) - monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) + monkeypatch.setattr("repositories.human_input.form_repository.naive_utc_now", lambda: fixed_now) ids = iter(["form-id", "del-web", "del-console", "del-backstage"]) - monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: next(ids)) + monkeypatch.setattr("repositories.human_input.form_repository.uuidv7", lambda: next(ids)) repo = HumanInputFormRepositoryImpl( tenant_id="tenant", @@ -850,6 +852,7 @@ def test_create_form_adds_console_and_backstage_recipients( workflow_execution_id="run", invoke_source="debugger", submission_actor_id="acc-1", + sessions=_sessions(repository_session), ) form_config = HumanInputNodeData( @@ -864,7 +867,7 @@ def test_create_form_adds_console_and_backstage_recipients( node_id="node", form_config=form_config, rendered_content="

hello

", - delivery_methods=[WebAppDeliveryMethod()], + delivery_methods=[InteractiveSurfaceDeliveryMethod()], display_in_ui=True, resolved_default_values={}, form_kind=HumanInputFormKind.RUNTIME, @@ -873,8 +876,6 @@ def test_create_form_adds_console_and_backstage_recipients( entity = repo.create_form(params) assert entity.id == "form-id" assert entity.expiration_time == fixed_now + timedelta(hours=form_config.timeout) - # Console token should take precedence when console recipient is present. - assert entity.submission_token is not None assert len(entity.recipients) == 3 repository_session.expire_all() persisted_form = repository_session.get(HumanInputForm, "form-id") @@ -890,8 +891,14 @@ def test_create_form_adds_console_and_backstage_recipients( } @staticmethod - def test_create_runtime_form_requires_workflow_even_when_conversation_is_present() -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant", app_id="app") + def test_create_runtime_form_requires_workflow_even_when_conversation_is_present( + repository_session: Session, + ) -> None: + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant", + app_id="app", + sessions=_sessions(repository_session), + ) params = FormCreateParams( workflow_execution_id=None, conversation_id="conversation", @@ -911,9 +918,44 @@ def test_create_runtime_form_requires_workflow_even_when_conversation_is_present with pytest.raises(ValueError, match="requires a workflow_execution_id"): repo.create_form(params) + params.workflow_execution_id = "run" + entity = repo.create_form(params) + + repository_session.expire_all() + persisted = repository_session.get(HumanInputForm, entity.id) + assert persisted is not None + assert persisted.workflow_run_id == "run" + assert persisted.conversation_id == "conversation" + + @staticmethod + def test_create_runtime_form_requires_workflow_without_conversation(repository_session: Session) -> None: + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant", + app_id="app", + sessions=_sessions(repository_session), + ) + params = FormCreateParams( + workflow_execution_id=None, + conversation_id=None, + node_id="node", + form_config=HumanInputNodeData( + title="Title", + delivery_methods=[], + form_content="hello", + inputs=[], + user_actions=[UserActionConfig(id="submit", title="Submit")], + ), + rendered_content="hello", + delivery_methods=[], + display_in_ui=True, + resolved_default_values={}, + ) + with pytest.raises(ValueError, match="requires a workflow_execution_id"): + repo.create_form(params) + @staticmethod def test_submission_get_by_token_returns_none_when_missing_or_form_missing(repository_session: Session) -> None: - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) assert repo.get_by_token("tok") is None _persist_recipient(repository_session, form_id="missing-form", access_token="orphan-token") @@ -928,17 +970,11 @@ def test_submission_repository_init_no_args() -> None: def test_submission_get_by_token_and_get_by_form_id_success_paths(repository_session: Session) -> None: form = _human_input_repository_persist_form(repository_session, form_id="f1", workflow_run_id=None) _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) record = repo.get_by_token("tok") assert record is not None assert record.access_token == "tok" - record = repo.get_by_form_id_and_recipient_type( - form_id=form.id, recipient_type=RecipientType.STANDALONE_WEB_APP - ) - assert record is not None - assert record.recipient_id == "r1" - record = repo.get_by_form_id(form.id) assert record is not None assert record.form_id == form.id @@ -946,8 +982,7 @@ def test_submission_get_by_token_and_get_by_form_id_success_paths(repository_ses @staticmethod def test_submission_get_by_form_id_returns_none_on_missing(repository_session: Session) -> None: - repo = HumanInputFormSubmissionRepository() - assert repo.get_by_form_id_and_recipient_type(form_id="f", recipient_type=RecipientType.CONSOLE) is None + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) assert repo.get_by_form_id("f") is None @staticmethod @@ -963,7 +998,7 @@ def record_statement(_connection, _cursor, statement: str, _parameters, _context engine = repository_session.get_bind() event.listen(engine, "before_cursor_execute", record_statement) try: - records = HumanInputFormSubmissionRepository().get_by_form_ids( + records = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)).get_by_form_ids( ["second", "missing", "first", "first"], tenant_id="tenant", app_id="app", workflow_run_id="run" ) assert len(statements) == 1 @@ -974,7 +1009,7 @@ def record_statement(_connection, _cursor, statement: str, _parameters, _context # Materialized records remain usable after the repository's session closes. assert len(statements) == 1 assert ( - HumanInputFormSubmissionRepository().get_by_form_ids( + HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)).get_by_form_ids( [], tenant_id="tenant", app_id="app", workflow_run_id="run" ) == {} @@ -988,9 +1023,9 @@ def test_mark_submitted_updates_and_raises_when_missing( repository_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: fixed_now = datetime(2024, 1, 1, 0, 0, 0) - monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) + monkeypatch.setattr("repositories.human_input.form_repository.naive_utc_now", lambda: fixed_now) - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) with pytest.raises(FormNotFoundError, match="form not found"): repo.mark_submitted( form_id="missing", @@ -1034,7 +1069,7 @@ def test_mark_submitted_does_not_overwrite_a_concurrent_stop_or_submission( form.status = status repository_session.commit() with pytest.raises(FormAlreadyHandledError) as error: - HumanInputFormSubmissionRepository().mark_submitted( + HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)).mark_submitted( form_id=form.id, recipient_id=None, selected_action_id="approve", @@ -1052,7 +1087,7 @@ def test_mark_submitted_serializes_select_and_file_payloads( repository_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: fixed_now = datetime(2024, 1, 1, 0, 0, 0) - monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) + monkeypatch.setattr("repositories.human_input.form_repository.naive_utc_now", lambda: fixed_now) form = _human_input_repository_persist_form(repository_session, form_id="f-complex", workflow_run_id=None) recipient = _persist_recipient( @@ -1093,7 +1128,7 @@ def test_mark_submitted_serializes_select_and_file_payloads( ], } - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) record = repo.mark_submitted( form_id=form.id, recipient_id=recipient.id, @@ -1112,7 +1147,7 @@ def test_mark_submitted_serializes_select_and_file_payloads( @staticmethod def test_mark_timeout_invalid_status_rolls_back(repository_session: Session) -> None: form = _human_input_repository_persist_form(repository_session, form_id="f", workflow_run_id=None) - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) with pytest.raises(_InvalidTimeoutStatusError, match="invalid timeout status"): repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.SUBMITTED) # type: ignore[arg-type] repository_session.expire_all() @@ -1128,7 +1163,7 @@ def test_mark_timeout_already_timed_out_returns_record(repository_session: Sessi workflow_run_id=None, status=HumanInputFormStatus.TIMEOUT, ) - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) record = repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.TIMEOUT, reason="r") assert record.status == HumanInputFormStatus.TIMEOUT @@ -1140,7 +1175,7 @@ def test_mark_timeout_submitted_raises_form_not_found(repository_session: Sessio workflow_run_id=None, status=HumanInputFormStatus.SUBMITTED, ) - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) with pytest.raises(FormNotFoundError, match="form already submitted"): repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.EXPIRED) @@ -1153,7 +1188,7 @@ def test_mark_timeout_updates_fields(repository_session: Session) -> None: form.submission_end_user_id = "eu" form.completed_by_recipient_id = "r" repository_session.commit() - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) record = repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.EXPIRED) repository_session.expire_all() persisted = repository_session.get(HumanInputForm, form.id) @@ -1168,7 +1203,7 @@ def test_mark_timeout_updates_fields(repository_session: Session) -> None: @staticmethod def test_mark_timeout_raises_when_form_missing(repository_session: Session) -> None: - repo = HumanInputFormSubmissionRepository() + repo = HumanInputFormSubmissionRepository(sessions=_sessions(repository_session)) with pytest.raises(FormNotFoundError, match="form not found"): repo.mark_timeout(form_id="missing", timeout_status=HumanInputFormStatus.TIMEOUT) @@ -1177,7 +1212,7 @@ def test_delivery_method_to_model_instant_message_creates_recipients( repository_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: "del-1") + monkeypatch.setattr("repositories.human_input.form_repository.uuidv7", lambda: "del-1") method = InstantMessageDeliveryMethod( config=InstantMessageDeliveryConfig( provider=InstantMessageProvider.SLACK, @@ -1280,7 +1315,7 @@ def test_build_email_recipients_uses_all_members_when_whole_workspace( @pytest.fixture -def repository_session(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: +def repository_session(sqlite_engine: Engine) -> Iterator[Session]: """Bind repository-owned sessions to an isolated SQLite database.""" tables = [ @@ -1292,9 +1327,5 @@ def repository_session(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) - ] TypeBase.metadata.create_all(sqlite_engine, tables=tables) repository_session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - monkeypatch.setattr( - "core.repositories.human_input_repository.session_factory.create_session", - repository_session_factory, - ) with repository_session_factory() as session: yield session diff --git a/api/tests/unit_tests/repositories/credentials/test_query_repository.py b/api/tests/unit_tests/repositories/credentials/test_query_repository.py index 2d1b2cf44e46df..8e598d521b6804 100644 --- a/api/tests/unit_tests/repositories/credentials/test_query_repository.py +++ b/api/tests/unit_tests/repositories/credentials/test_query_repository.py @@ -8,7 +8,7 @@ from uuid import uuid4 import pytest -from sqlalchemy import delete, select, update +from sqlalchemy import Engine, delete, event, select, update from sqlalchemy.orm import Session, sessionmaker from core.plugin.entities.plugin_daemon import CredentialType as TriggerCredentialType @@ -45,6 +45,58 @@ def repository(sqlite_session_factory: sessionmaker[Session]) -> CredentialQuery return CredentialQueryRepository(session_factory=sqlite_session_factory) +@pytest.mark.parametrize("has_default", [False, True]) +def test_default_tool_id_is_a_scoped_single_column_query( + sqlite_session: Session, + sqlite_engine: Engine, + repository: CredentialQueryRepository, + has_default: bool, +) -> None: + for credential_id, tenant, provider, day, default in [ + ("oldest", "workspace", "tool", 2, False), + ("default", "workspace", "tool", 3, has_default), + ("later-default", "workspace", "tool", 4, has_default), + ("foreign", "other-workspace", "tool", 1, True), + ("other-provider", "workspace", "other-tool", 1, True), + ]: + row = BuiltinToolProvider( + name=credential_id, + tenant_id=tenant, + user_id="owner", + provider=provider, + is_default=default, + visibility=PermissionEnum.PARTIAL_TEAM, + # Selecting an ID must neither parse credentials nor fetch member permissions. + encrypted_credentials="not credential JSON", + ) + row.id = credential_id + row.created_at = datetime(2026, 1, day) + sqlite_session.add(row) + sqlite_session.commit() + statements: list[str] = [] + + def record( + _connection: object, + _cursor: object, + statement: str, + _parameters: object, + _context: object, + _executemany: bool, + ) -> None: + statements.append(statement.lower()) + + event.listen(sqlite_engine, "before_cursor_execute", record) + try: + result = repository.default_tool_credential_id(workspace_id="workspace", provider="tool") + finally: + event.remove(sqlite_engine, "before_cursor_execute", record) + assert result == ("default" if has_default else "oldest") + assert len(statements) == 1 + table = BuiltinToolProvider.__tablename__ + assert statements[0].split("\nfrom")[0].strip() == f"select {table}.id" + assert "limit" in statements[0] + + @pytest.mark.parametrize("provider", ["openai", "langgenius/openai/openai"]) def test_model_aliases_preserve_visibility_and_tenant_scope( sqlite_session: Session, repository: CredentialQueryRepository, provider: str @@ -293,6 +345,7 @@ def test_empty_lists_do_not_require_related_rows(repository: CredentialQueryRepo repository.list_datasources(workspace_id="empty", provider="notion", plugin_id="plugin", actor_id="actor") == [] ) assert repository.list_tools(workspace_id="empty", provider="tool", actor_id="actor") == [] + assert repository.default_tool_credential_id(workspace_id="empty", provider="tool") is None assert repository.list_trigger_subscriptions(workspace_id="empty", provider="trigger", actor_id="actor") == [] diff --git a/api/tests/unit_tests/repositories/workflow/test_imports.py b/api/tests/unit_tests/repositories/workflow/test_imports.py new file mode 100644 index 00000000000000..96392a96710d5c --- /dev/null +++ b/api/tests/unit_tests/repositories/workflow/test_imports.py @@ -0,0 +1,26 @@ +"""Repositories must initialize without application services being imported first.""" + +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize( + "module", + [ + "repositories.workflow.definition_repository", + "repositories.agent.creation_repository", + "repositories.agent.workflow_binding_repository", + ], +) +def test_repository_imports_in_a_fresh_process(module: str) -> None: + result = subprocess.run( + [sys.executable, "-c", f"import {module}"], + cwd=Path(__file__).resolve().parents[4], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr diff --git a/api/tests/unit_tests/repositories/workflow/test_node_execution_repository.py b/api/tests/unit_tests/repositories/workflow/test_node_execution_repository.py new file mode 100644 index 00000000000000..8979ae04417fa4 --- /dev/null +++ b/api/tests/unit_tests/repositories/workflow/test_node_execution_repository.py @@ -0,0 +1,63 @@ +"""Execution response reads materialize creator data within its owner scope.""" + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from models.account import TenantAccountJoin, TenantAccountRole +from models.enums import CreatorUserRole +from models.workflow import WorkflowNodeExecutionModel +from repositories.workflow.node_execution_repository import WorkflowNodeExecutionRepository +from tests.helpers.unit.model_factories import make_account, make_end_user, make_tenant + + +@pytest.mark.parametrize("tenant_id", ["tenant-1", "other-tenant"]) +def test_account_response_requires_execution_tenant_membership( + sqlite_session_factory: sessionmaker[Session], tenant_id: str +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + make_account(), + make_tenant(tenant_id=tenant_id), + TenantAccountJoin(tenant_id=tenant_id, account_id="account-1", role=TenantAccountRole.OWNER), + ] + ) + execution = WorkflowNodeExecutionModel( + tenant_id="tenant-1", + app_id="app-1", + created_by="account-1", + created_by_role=CreatorUserRole.ACCOUNT, + ) + execution.offload_data = [] + result = WorkflowNodeExecutionRepository(sqlite_session_factory).execution_record(execution, include_details=True) + assert result["created_by_end_user"] is None + if tenant_id == "tenant-1": + assert result["created_by_account"]["id"] == "account-1" + assert isinstance(result["created_by_account"], dict) + else: + assert result["created_by_account"] is None + + +@pytest.mark.parametrize( + ("tenant_id", "app_id"), [("tenant-1", "app-1"), ("other-tenant", "app-1"), ("tenant-1", "other-app")] +) +def test_end_user_response_requires_execution_tenant_and_app( + sqlite_session_factory: sessionmaker[Session], tenant_id: str, app_id: str +) -> None: + with sqlite_session_factory.begin() as session: + session.add(make_end_user(tenant_id=tenant_id, app_id=app_id)) + execution = WorkflowNodeExecutionModel( + tenant_id="tenant-1", + app_id="app-1", + created_by="end-user-1", + created_by_role=CreatorUserRole.END_USER, + ) + execution.offload_data = [] + result = WorkflowNodeExecutionRepository(sqlite_session_factory).execution_record(execution, include_details=True) + assert result["created_by_account"] is None + if (tenant_id, app_id) == ("tenant-1", "app-1"): + assert result["created_by_end_user"]["id"] == "end-user-1" + assert result["created_by_end_user"]["session_id"] == "session-1" + assert isinstance(result["created_by_end_user"], dict) + else: + assert result["created_by_end_user"] is None diff --git a/api/tests/unit_tests/services/agent/test_provider_service.py b/api/tests/unit_tests/services/agent/test_provider_service.py new file mode 100644 index 00000000000000..076e31e023c2f7 --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_provider_service.py @@ -0,0 +1,41 @@ +"""Agent provider delegation has no database dependencies.""" + +from dataclasses import dataclass + +import pytest + +from core.plugin.impl.exc import PluginDaemonClientSideError +from services import agent_service + + +@dataclass +class Provider: + plugin_id: str + + +class PluginClient: + def __init__(self) -> None: + self.provider: Provider = Provider("plugin") + self.calls: list[tuple[str, str | None]] = [] + self.failure: PluginDaemonClientSideError | None = None + + def fetch_agent_strategy_providers(self, tenant_id: str) -> list[Provider]: + self.calls.append((tenant_id, None)) + return [self.provider] + + def fetch_agent_strategy_provider(self, tenant_id: str, name: str) -> Provider: + self.calls.append((tenant_id, name)) + if self.failure: + raise self.failure + return self.provider + + +def test_agent_provider_queries_preserve_scope_and_errors(monkeypatch: pytest.MonkeyPatch) -> None: + client = PluginClient() + monkeypatch.setattr(agent_service, "PluginAgentClient", lambda: client) + assert agent_service.AgentService.list_agent_providers("user", "tenant") == [client.provider] + assert agent_service.AgentService.get_agent_provider("user", "tenant", "name") is client.provider + client.failure = PluginDaemonClientSideError("Plugin not found") + with pytest.raises(ValueError, match="Plugin not found"): + agent_service.AgentService.get_agent_provider("user", "tenant", "missing") + assert client.calls == [("tenant", None), ("tenant", "name"), ("tenant", "missing")] diff --git a/api/tests/unit_tests/services/agent/test_roster_service.py b/api/tests/unit_tests/services/agent/test_roster_service.py index 613cff846cafdc..067f92f7cf09ae 100644 --- a/api/tests/unit_tests/services/agent/test_roster_service.py +++ b/api/tests/unit_tests/services/agent/test_roster_service.py @@ -535,7 +535,7 @@ def test_save_workflow_composer_dispatches_save_strategy(monkeypatch, strategy, monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_draft_save_payload", lambda payload: None) monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: _workflow()) - monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: None) + monkeypatch.setattr(AgentComposerService, "_lock_workflow_binding", lambda **kwargs: None) monkeypatch.setattr( AgentComposerService, "_get_agent_if_present", @@ -608,7 +608,7 @@ def test_save_workflow_composer_commits_before_retiring_replaced_inline_agent( current_snapshot_id="version-1", ) monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **_kwargs: _workflow()) - monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **_kwargs: old_binding) + monkeypatch.setattr(AgentComposerService, "_lock_workflow_binding", lambda **_kwargs: old_binding) monkeypatch.setattr(AgentComposerService, "_save_as_new_agent", lambda **_kwargs: new_binding) monkeypatch.setattr( AgentComposerService, @@ -3224,7 +3224,7 @@ def test_copy_workflow_composer_from_roster_creates_inline_agent_and_preserves_n captured: dict[str, object] = {} monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: workflow) - monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: binding) + monkeypatch.setattr(AgentComposerService, "_lock_workflow_binding", lambda **kwargs: binding) monkeypatch.setattr(AgentComposerService, "_require_agent", lambda **kwargs: roster_agent) monkeypatch.setattr(AgentComposerService, "_require_version", lambda **kwargs: source_version) @@ -3273,7 +3273,7 @@ def test_copy_workflow_composer_from_roster_rejects_stale_source_snapshot( monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: _workflow()) monkeypatch.setattr( AgentComposerService, - "_get_workflow_binding", + "_lock_workflow_binding", lambda **kwargs: WorkflowAgentNodeBinding( tenant_id="tenant-1", app_id="app-1", @@ -3353,7 +3353,7 @@ def test_copy_workflow_composer_from_roster_rejects_unpublished_source( active_config_snapshot_id="roster-version-1", ) monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: _workflow()) - monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: binding) + monkeypatch.setattr(AgentComposerService, "_lock_workflow_binding", lambda **kwargs: binding) monkeypatch.setattr(AgentComposerService, "_require_agent", lambda **kwargs: source_agent) require_version = MagicMock(side_effect=AssertionError("unpublished source must fail before loading snapshot")) monkeypatch.setattr(AgentComposerService, "_require_version", require_version) @@ -3406,7 +3406,7 @@ def test_copy_workflow_composer_from_roster_is_idempotent_when_already_inline( serialize_calls = [] session = sqlite_session monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: _workflow()) - monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: inline_binding) + monkeypatch.setattr(AgentComposerService, "_lock_workflow_binding", lambda **kwargs: inline_binding) monkeypatch.setattr(AgentComposerService, "_get_agent_if_present", lambda **kwargs: inline_agent) monkeypatch.setattr(AgentComposerService, "_get_version_if_present", lambda **kwargs: inline_version) @@ -3494,7 +3494,7 @@ def test_copy_workflow_composer_from_roster_rejects_invalid_source_binding( active_config_snapshot_id="version-1", ) monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: _workflow()) - monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: binding) + monkeypatch.setattr(AgentComposerService, "_lock_workflow_binding", lambda **kwargs: binding) monkeypatch.setattr(AgentComposerService, "_require_agent", lambda **kwargs: source_agent) with pytest.raises(InvalidComposerConfigError, match=expected_message): diff --git a/api/tests/unit_tests/services/agent/test_tool_invocation_service.py b/api/tests/unit_tests/services/agent/test_tool_invocation_service.py new file mode 100644 index 00000000000000..2e56a28e2e48b1 --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_tool_invocation_service.py @@ -0,0 +1,278 @@ +"""Unit tests for the Agent tool inner invoke service with SQLite-backed app lookup.""" + +from collections.abc import Generator +from dataclasses import dataclass, field, replace +from unittest.mock import create_autospec, patch + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.tools.__base.tool_runtime import ToolRuntime +from core.tools.entities.tool_entities import ToolInvokeFrom, ToolInvokeMessage, ToolProviderType +from core.tools.errors import ( + ToolInvokeError, + ToolParameterValidationError, + ToolProviderCredentialValidationError, + ToolProviderNotFoundError, +) +from extensions.application_services.workflow import WorkflowExecutionDependencies +from models.enums import AppStatus +from models.model import App, AppMode +from repositories.app.agent_app_repository import AgentAppRepository +from services.agent.tool_invocation_service import AgentToolInnerService +from services.credentials.refresh import CredentialRefreshUnavailableError +from services.entities.agent_tool_inner import AgentToolInvokeRequest +from services.errors.agent_tool_inner import AgentToolInnerServiceError +from services.tools.agent_invocation_gateway import AgentToolInvocationGateway +from services.workflow.execution.ports import WorkflowRuntime, WorkflowToolInvoker +from services.workflow.variable_contracts import WorkflowExecutionVariables + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +OTHER_TENANT_ID = "22222222-2222-2222-2222-222222222222" +USER_ID = "33333333-3333-3333-3333-333333333333" +APP_ID = "44444444-4444-4444-4444-444444444444" + + +def _persist_app(sqlite_session: Session, *, tenant_id: str = TENANT_ID) -> App: + app = App( + id=APP_ID, + tenant_id=tenant_id, + name="Test App", + description="", + mode=AppMode.CHAT, + status=AppStatus.NORMAL, + enable_site=False, + enable_api=False, + max_active_requests=None, + ) + sqlite_session.add(app) + sqlite_session.commit() + sqlite_session.expunge_all() + return app + + +def _request() -> AgentToolInvokeRequest: + return AgentToolInvokeRequest.model_validate( + { + "caller": { + "tenant_id": TENANT_ID, + "user_id": USER_ID, + "user_from": "account", + "app_id": APP_ID, + "invoke_from": "service-api", + "conversation_id": "conversation-1", + "workflow_id": "workflow-1", + "workflow_run_id": "workflow-run-1", + "node_id": "node-1", + "node_execution_id": "node-exec-1", + "agent_id": "agent-1", + "agent_config_version_id": "snapshot-1", + "trace_session_id": "trace-session-1", + }, + "tool": { + "provider_type": "plugin", + "provider_id": "langgenius/search/search", + "tool_name": "search", + "credential_id": "credential-1", + "tool_parameters": {"query": "dify"}, + "runtime_parameters": {"region": "us"}, + }, + } + ) + + +def _messages() -> Generator[ToolInvokeMessage]: + yield ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.TEXT, + message=ToolInvokeMessage.TextMessage(text="ok"), + ) + + +@dataclass +class _FakeTool: + """Resolved tool stub whose runtime metadata is populated by the gateway.""" + + runtime: ToolRuntime = field(default_factory=lambda: ToolRuntime(tenant_id=TENANT_ID)) + + +def _tool() -> _FakeTool: + return _FakeTool() + + +def _variables() -> WorkflowExecutionVariables: + return create_autospec(WorkflowExecutionVariables, instance=True, spec_set=True) + + +def _service( + session: Session, + runtime: WorkflowRuntime, + variables: WorkflowExecutionVariables | None = None, +) -> AgentToolInnerService: + return AgentToolInnerService( + apps=AgentAppRepository(session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False)), + tools=AgentToolInvocationGateway( + variables=variables if variables is not None else _variables(), + runtime=runtime, + ), + ) + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +@pytest.mark.parametrize( + ("workflow_id", "expected_tool_invoke_from"), + [ + ("workflow-1", ToolInvokeFrom.WORKFLOW), + (None, ToolInvokeFrom.AGENT), + ], +) +def test_invoke_uses_agent_tool_runtime_and_returns_observation( + sqlite_session: Session, + workflow_id: str | None, + expected_tool_invoke_from: ToolInvokeFrom, + *, + workflow_runtime: WorkflowExecutionDependencies, +) -> None: + fake_tool = _tool() + variables = _variables() + tool_invoker = create_autospec(WorkflowToolInvoker, instance=True, spec_set=True) + tool_invoker.return_value = _messages() + runtime = replace(workflow_runtime, tool_invoker=tool_invoker) + _persist_app(sqlite_session) + + with ( + patch( + "services.tools.agent_invocation_gateway.ToolManager.get_agent_tool_runtime", + return_value=fake_tool, + ) as mock_get_runtime, + patch( + "services.tools.agent_invocation_gateway.ToolFileMessageTransformer.transform_tool_invoke_messages", + side_effect=lambda messages, **_kwargs: messages, + ), + ): + request = _request() + request.caller.workflow_id = workflow_id + response = _service(sqlite_session, runtime, variables).invoke(request) + + assert response.observation == "ok" + assert response.metadata == { + "provider_type": "plugin", + "provider_id": "langgenius/search/search", + "tool_name": "search", + } + agent_tool = mock_get_runtime.call_args.kwargs["agent_tool"] + assert mock_get_runtime.call_args.kwargs["draft_variable_saver"] is variables.saver_factory + assert agent_tool.provider_type is ToolProviderType.PLUGIN + assert agent_tool.tool_parameters == {"region": "us"} + assert fake_tool.runtime.dify_run_context is not None + assert fake_tool.runtime.dify_run_context.tenant_id == TENANT_ID + assert fake_tool.runtime.dify_run_context.app_id == APP_ID + assert fake_tool.runtime.dify_run_context.user_id == USER_ID + assert fake_tool.runtime.dify_run_context.user_from is UserFrom.ACCOUNT + assert fake_tool.runtime.dify_run_context.invoke_from is InvokeFrom.SERVICE_API + assert fake_tool.runtime.dify_run_context.trace_session_id == "trace-session-1" + assert fake_tool.runtime.tool_invoke_from is expected_tool_invoke_from + tool_invoker.assert_called_once() + assert "session" not in tool_invoker.call_args.kwargs + assert not sqlite_session.in_transaction() + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_invoke_raises_app_not_found(sqlite_session: Session, *, workflow_runtime: WorkflowRuntime) -> None: + with pytest.raises(AgentToolInnerServiceError) as exc_info: + _service(sqlite_session, workflow_runtime).invoke(_request()) + + assert exc_info.value.error_code == "app_not_found" + assert exc_info.value.status_code == 404 + assert exc_info.value.description == "App not found." + assert not sqlite_session.in_transaction() + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_invoke_raises_app_tenant_mismatch_when_app_belongs_to_other_tenant( + sqlite_session: Session, *, workflow_runtime: WorkflowRuntime +) -> None: + _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) + + with pytest.raises(AgentToolInnerServiceError) as exc_info: + _service(sqlite_session, workflow_runtime).invoke(_request()) + + assert exc_info.value.error_code == "app_tenant_mismatch" + assert exc_info.value.status_code == 403 + assert exc_info.value.description == "App does not belong to the caller tenant." + assert not sqlite_session.in_transaction() + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_invoke_maps_tool_runtime_app_not_found_value_error_to_specific_error_code( + sqlite_session: Session, *, workflow_runtime: WorkflowExecutionDependencies +) -> None: + fake_tool = _tool() + tool_invoker = create_autospec(WorkflowToolInvoker, instance=True, spec_set=True) + tool_invoker.side_effect = ValueError("app not found") + runtime = replace(workflow_runtime, tool_invoker=tool_invoker) + _persist_app(sqlite_session) + + with patch("services.tools.agent_invocation_gateway.ToolManager.get_agent_tool_runtime", return_value=fake_tool): + with pytest.raises(AgentToolInnerServiceError) as exc_info: + _service(sqlite_session, runtime).invoke(_request()) + + assert exc_info.value.error_code == "app_not_found" + assert exc_info.value.status_code == 404 + assert exc_info.value.description == "App not found." + assert not sqlite_session.in_transaction() + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_invoke_maps_tool_invoke_error_without_private_tool_engine_helper( + sqlite_session: Session, *, workflow_runtime: WorkflowExecutionDependencies +) -> None: + fake_tool = _tool() + tool_invoker = create_autospec(WorkflowToolInvoker, instance=True, spec_set=True) + tool_invoker.side_effect = ToolInvokeError("workflow crashed") + runtime = replace(workflow_runtime, tool_invoker=tool_invoker) + _persist_app(sqlite_session) + + with patch("services.tools.agent_invocation_gateway.ToolManager.get_agent_tool_runtime", return_value=fake_tool): + with pytest.raises(AgentToolInnerServiceError) as exc_info: + _service(sqlite_session, runtime).invoke(_request()) + + assert exc_info.value.error_code == "agent_tool_invoke_failed" + assert not sqlite_session.in_transaction() + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_invoke_preserves_credential_refresh_unavailable_error( + sqlite_session: Session, *, workflow_runtime: WorkflowRuntime +) -> None: + _persist_app(sqlite_session) + error = CredentialRefreshUnavailableError() + + with patch("services.tools.agent_invocation_gateway.ToolManager.get_agent_tool_runtime", side_effect=error): + with pytest.raises(CredentialRefreshUnavailableError) as exc_info: + _service(sqlite_session, workflow_runtime).invoke(_request()) + + assert exc_info.value is error + assert not sqlite_session.in_transaction() + + +@pytest.mark.parametrize( + ("error", "expected_code"), + [ + (ToolProviderNotFoundError("provider gone"), "agent_tool_declaration_not_found"), + (ToolProviderCredentialValidationError("credential invalid"), "agent_tool_credential_invalid"), + (ToolParameterValidationError("query is required"), "tool_parameters_invalid"), + ], +) +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_invoke_maps_runtime_lookup_errors_to_service_error_codes( + error: Exception, expected_code: str, sqlite_session: Session, *, workflow_runtime: WorkflowRuntime +) -> None: + _persist_app(sqlite_session) + + with patch("services.tools.agent_invocation_gateway.ToolManager.get_agent_tool_runtime", side_effect=error): + with pytest.raises(AgentToolInnerServiceError) as exc_info: + _service(sqlite_session, workflow_runtime).invoke(_request()) + + assert exc_info.value.error_code == expected_code + assert not sqlite_session.in_transaction() diff --git a/api/tests/unit_tests/services/tools/test_plugin_invocation.py b/api/tests/unit_tests/services/tools/test_plugin_invocation.py new file mode 100644 index 00000000000000..13979192abce55 --- /dev/null +++ b/api/tests/unit_tests/services/tools/test_plugin_invocation.py @@ -0,0 +1,150 @@ +from collections.abc import Generator +from typing import override +from unittest.mock import patch + +import pytest +from sqlalchemy import Engine +from sqlalchemy.orm import Session, sessionmaker + +from core.app.apps.draft_variable_saver import DraftVariableSaverFactory +from core.tools.__base.tool import Tool +from core.tools.__base.tool_runtime import ToolRuntime +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage, ToolProviderType +from extensions.application_services.workflow import ( + WorkflowExecutionDependencies, + build_workflow_execution_dependencies, +) +from models import Account +from services.tools.plugin_invocation import PluginToolBackwardsInvocation + + +class _TrackingSession(Session): + was_closed: bool = False + + @override + def close(self) -> None: + self.was_closed = True + super().close() + + +class _LazyTool(Tool): + def __init__(self) -> None: + super().__init__( + entity=ToolEntity( + identity=ToolIdentity( + author="test", + name="lazy-plugin-tool", + label=I18nObject(en_US="Lazy plugin tool"), + provider="test-provider", + ) + ), + runtime=ToolRuntime(tenant_id="tenant-1", user_id="user-1"), + ) + self.sessions: list[_TrackingSession] = [] + + @override + def tool_provider_type(self) -> ToolProviderType: + return ToolProviderType.PLUGIN + + @override + def _invoke( + self, + session: Session, + user_id: str, + tool_parameters: dict[str, object], + conversation_id: str | None = None, + app_id: str | None = None, + message_id: str | None = None, + ) -> Generator[ToolInvokeMessage]: + assert isinstance(session, _TrackingSession) + assert user_id == "user-1" + assert tool_parameters == {"query": "hello"} + self.sessions.append(session) + + assert not session.was_closed + yield self.create_text_message("first") + assert not session.was_closed + yield self.create_text_message("second") + + +def _runtime(sqlite_engine: Engine) -> WorkflowExecutionDependencies: + sessions: sessionmaker[Session] = sessionmaker( + bind=sqlite_engine, + class_=_TrackingSession, + expire_on_commit=False, + ) + return build_workflow_execution_dependencies(sessions) + + +def _unused_draft_variable_saver(_tenant_id: str, _account: Account) -> DraftVariableSaverFactory: + raise AssertionError("the patched tool lookup must not call the draft variable saver") + + +def _invoke(runtime: WorkflowExecutionDependencies) -> Generator[ToolInvokeMessage]: + return PluginToolBackwardsInvocation.invoke_tool( + tenant_id="tenant-1", + user_id="user-1", + tool_type=ToolProviderType.PLUGIN, + provider="test-provider", + tool_name="lazy-plugin-tool", + tool_parameters={"query": "hello"}, + workflow_runtime=runtime, + draft_variable_saver=_unused_draft_variable_saver, + ) + + +def test_plugin_invocation_closes_transform_and_session_when_consumer_stops_early(sqlite_engine: Engine) -> None: + runtime = _runtime(sqlite_engine) + tool = _LazyTool() + transform_closed = False + + def tracking_transform(messages: Generator[ToolInvokeMessage], **_kwargs: object) -> Generator[ToolInvokeMessage]: + nonlocal transform_closed + try: + yield from messages + finally: + transform_closed = True + + with ( + patch( + "services.tools.plugin_invocation.ToolManager.get_tool_runtime_from_plugin", return_value=tool + ) as get_runtime, + patch( + "services.tools.plugin_invocation.ToolFileMessageTransformer.transform_tool_invoke_messages", + side_effect=tracking_transform, + ), + ): + messages = _invoke(runtime) + get_runtime.assert_called_once() + assert tool.sessions == [] + + assert next(messages).message == ToolInvokeMessage.TextMessage(text="first") + assert not tool.sessions[0].was_closed + messages.close() + + assert transform_closed + assert tool.sessions[0].was_closed + + +def test_plugin_invocation_closes_session_when_transformer_raises(sqlite_engine: Engine) -> None: + runtime = _runtime(sqlite_engine) + tool = _LazyTool() + + def failing_transform(messages: Generator[ToolInvokeMessage], **_kwargs: object) -> Generator[ToolInvokeMessage]: + yield next(messages) + raise RuntimeError("transform failed") + + with ( + patch("services.tools.plugin_invocation.ToolManager.get_tool_runtime_from_plugin", return_value=tool), + patch( + "services.tools.plugin_invocation.ToolFileMessageTransformer.transform_tool_invoke_messages", + side_effect=failing_transform, + ), + ): + messages = _invoke(runtime) + assert next(messages).message == ToolInvokeMessage.TextMessage(text="first") + with pytest.raises(RuntimeError, match="transform failed"): + next(messages) + + assert tool.sessions[0].was_closed