From a048f350994e49897a30fecf71c161335c943637 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?= Date: Tue, 28 Jul 2026 21:16:05 +0800 Subject: [PATCH] refactor: introduce agent working environment architecture (#39480) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .github/workflows/web-e2e.yml | 1 + api/clients/agent_backend/__init__.py | 12 - api/clients/agent_backend/request_builder.py | 106 +- api/clients/agent_backend/session_cleanup.py | 100 - api/configs/extra/agent_backend_config.py | 5 +- api/controllers/console/agent/roster.py | 36 +- .../console/app/agent_app_sandbox.py | 76 +- api/controllers/console/app/completion.py | 33 +- api/controllers/console/app/workflow.py | 14 +- .../console/snippets/snippet_workflow.py | 15 +- api/core/app/apps/advanced_chat/app_runner.py | 16 +- api/core/app/apps/agent_app/app_generator.py | 205 +- api/core/app/apps/agent_app/app_runner.py | 188 +- .../apps/agent_app/runtime_request_builder.py | 13 +- api/core/app/apps/agent_app/session_store.py | 339 ++- api/core/app/apps/workflow/app_runner.py | 17 +- api/core/app/entities/app_invoke_entities.py | 13 +- api/core/app/workflow/layers/persistence.py | 23 +- ...lery_workflow_node_execution_repository.py | 22 + api/core/repositories/factory.py | 2 + ...hemy_workflow_node_execution_repository.py | 31 +- .../workflow/node_execution_process_data.py | 27 + api/core/workflow/node_factory.py | 4 +- .../workflow/nodes/agent_v2/agent_node.py | 304 ++- .../nodes/agent_v2/binding_resolver.py | 47 +- .../nodes/agent_v2/runtime_request_builder.py | 11 +- .../nodes/agent_v2/session_cleanup_layer.py | 126 -- .../workflow/nodes/agent_v2/session_store.py | 373 ++-- .../agent_v2/workspace_retirement_layer.py | 89 + api/extensions/ext_celery.py | 1 + ...tore_workflow_node_execution_repository.py | 6 + ...536b3feb_add_agent_home_snapshot_ledger.py | 64 + ...57_replace_agent_runtime_sessions_with_.py | 153 ++ api/models/__init__.py | 22 +- api/models/agent.py | 218 +- api/models/model.py | 8 + api/models/workflow.py | 1 + api/openapi/markdown/console-openapi.md | 31 +- api/services/agent/composer_service.py | 397 +++- api/services/agent/dsl_service.py | 17 +- api/services/agent/errors.py | 6 + api/services/agent/home_snapshot_service.py | 238 +++ api/services/agent/retirement_service.py | 166 ++ api/services/agent/roster_service.py | 429 ++-- .../agent/workflow_publish_service.py | 65 +- api/services/agent/workspace_service.py | 469 +++++ api/services/agent_app_sandbox_service.py | 452 ++-- api/services/app_dsl_service.py | 15 +- api/services/app_service.py | 74 +- api/services/conversation_service.py | 106 +- api/services/data_migration/import_service.py | 14 +- api/services/snippet_dsl_service.py | 28 +- api/services/snippet_service.py | 39 +- api/services/workflow_service.py | 33 +- .../agent_backend_session_cleanup_task.py | 72 - .../app_generate/resume_agent_app_task.py | 1 + api/tasks/collect_agent_resources_task.py | 75 + api/tasks/remove_app_and_related_data_task.py | 151 -- api/tasks/workflow_node_execution_tasks.py | 4 +- .../services/test_app_dsl_service.py | 9 +- .../services/test_app_service.py | 5 + .../services/test_conversation_service.py | 162 +- .../services/test_workflow_service.py | 3 +- ...anup_composition_compositor_integration.py | 134 -- .../clients/agent_backend/test_client.py | 1 + .../clients/agent_backend/test_fake_client.py | 1 + .../agent_backend/test_request_builder.py | 107 +- .../agent_backend/test_session_cleanup.py | 124 -- .../common/test_agent_app_parameters.py | 1 + .../console/agent/test_agent_controllers.py | 78 +- .../console/app/test_agent_app_sandbox.py | 141 +- .../unit_tests/controllers/test_swagger.py | 10 +- .../core/agent/test_publish_visibility.py | 2 + .../app/apps/agent_app/test_app_generator.py | 136 +- .../app/apps/agent_app/test_app_runner.py | 254 +-- .../app/apps/agent_app/test_resolve_agent.py | 180 +- .../agent_app/test_runtime_request_builder.py | 27 +- .../app/apps/agent_app/test_session_store.py | 415 +--- .../app/workflow/layers/test_persistence.py | 2 + .../app/workflow/test_persistence_layer.py | 61 +- ...lery_workflow_node_execution_repository.py | 17 + ...hemy_workflow_node_execution_repository.py | 35 +- ...rkflow_node_execution_conflict_handling.py | 1 + .../nodes/agent_v2/test_agent_node.py | 387 ++-- .../nodes/agent_v2/test_binding_resolver.py | 72 +- .../agent_v2/test_runtime_request_builder.py | 18 +- .../agent_v2/test_session_cleanup_layer.py | 256 --- .../nodes/agent_v2/test_session_store.py | 624 +++--- .../test_workspace_retirement_layer.py | 75 + .../events/test_app_event_signals.py | 3 +- ..._api_workflow_node_execution_repository.py | 28 +- ...tore_workflow_node_execution_repository.py | 35 + .../unit_tests/extensions/test_celery_ssl.py | 3 +- .../test_agent_workspace_migration.py | 140 ++ .../services/agent/test_agent_dsl_service.py | 20 +- .../services/agent/test_agent_services.py | 1830 ++++++++++++++--- .../agent/test_home_snapshot_service.py | 181 ++ .../services/agent/test_retirement_service.py | 205 ++ .../agent/test_workflow_publish_service.py | 125 +- .../services/agent/test_workspace_service.py | 422 ++++ .../test_agent_app_sandbox_service.py | 631 ++---- .../unit_tests/services/test_app_service.py | 73 +- .../services/test_conversation_service.py | 56 +- .../services/test_snippet_dsl_service.py | 17 + .../services/test_snippet_service.py | 13 +- .../services/test_workflow_service.py | 3 +- ...test_agent_backend_session_cleanup_task.py | 51 - .../test_collect_agent_resources_task.py | 81 + .../test_remove_app_and_related_data_task.py | 248 +-- .../tasks/test_resume_agent_app_task.py | 1 + .../test_workflow_node_execution_identity.py | 28 + api/uv.lock | 1 + dify-agent/.example.env | 57 +- dify-agent/.gitignore | 1 + .../concepts/runtime-resources/index.md | 194 ++ .../docs/dify-agent/get-started/index.md | 32 +- dify-agent/docs/dify-agent/guide/index.md | 175 +- .../user-manual/shell-layer/index.md | 305 ++- dify-agent/mkdocs.yml | 1 + dify-agent/pyproject.toml | 2 + .../src/dify_agent/adapters/shell/__init__.py | 14 - .../src/dify_agent/adapters/shell/config.py | 61 - .../adapters/shell/enterprise/__init__.py | 3 - .../adapters/shell/enterprise/enterprise.py | 235 --- .../src/dify_agent/adapters/shell/factory.py | 31 - .../dify_agent/adapters/shell/protocols.py | 54 +- .../src/dify_agent/adapters/shell/shellctl.py | 429 +++- .../agent_stub/server/agent_stub_files.py | 7 +- dify-agent/src/dify_agent/client/_client.py | 178 +- .../src/dify_agent/layers/runtime/__init__.py | 4 + .../src/dify_agent/layers/runtime/configs.py | 17 + .../src/dify_agent/layers/runtime/layer.py | 75 + .../src/dify_agent/layers/shell/__init__.py | 8 +- .../src/dify_agent/layers/shell/configs.py | 31 +- .../src/dify_agent/layers/shell/layer.py | 358 +--- .../src/dify_agent/protocol/__init__.py | 62 +- .../dify_agent/protocol/execution_binding.py | 44 + .../src/dify_agent/protocol/home_snapshot.py | 42 + dify-agent/src/dify_agent/protocol/sandbox.py | 254 --- dify-agent/src/dify_agent/protocol/schemas.py | 10 +- .../src/dify_agent/protocol/workspace.py | 84 + .../dify_agent/runtime/compositor_factory.py | 44 +- dify-agent/src/dify_agent/runtime/runner.py | 58 +- .../dify_agent/runtime_backend/__init__.py | 63 + .../src/dify_agent/runtime_backend/e2b.py | 405 ++++ .../dify_agent/runtime_backend/enterprise.py | 234 +++ .../src/dify_agent/runtime_backend/errors.py | 73 + .../src/dify_agent/runtime_backend/leases.py | 35 + .../src/dify_agent/runtime_backend/local.py | 331 +++ .../src/dify_agent/runtime_backend/profile.py | 137 ++ .../dify_agent/runtime_backend/protocols.py | 228 ++ .../dify_agent/runtime_backend/shellctl.py | 156 ++ dify-agent/src/dify_agent/server/app.py | 48 +- .../dify_agent/server/execution_bindings.py | 73 + .../src/dify_agent/server/home_snapshots.py | 83 + .../server/routes/execution_bindings.py | 55 + .../server/routes/home_snapshots.py | 72 + .../dify_agent/server/routes/sandbox_files.py | 73 - .../server/routes/workspace_files.py | 67 + .../src/dify_agent/server/sandbox_files.py | 428 ---- dify-agent/src/dify_agent/server/settings.py | 80 +- .../src/dify_agent/server/workspace_files.py | 222 ++ .../runtime_backend/run_local_integration.sh | 40 + .../test_working_environment.py | 221 ++ .../dify_agent/adapters/shell/test_config.py | 54 +- .../adapters/shell/test_enterprise.py | 42 - .../adapters/shell/test_shellctl.py | 370 +++- .../server/test_agent_stub_files.py | 20 + .../local/dify_agent/client/test_client.py | 248 ++- .../dify_agent/layers/config/test_layer.py | 6 - .../dify_agent/layers/drive/test_layer.py | 6 - .../dify_agent/layers/shell/test_configs.py | 6 - .../dify_agent/layers/shell/test_layer.py | 457 ++-- .../dify_agent/layers/test_runtime_layer.py | 61 + .../protocol/test_sandbox_locator.py | 178 -- .../protocol/test_working_environment.py | 47 + .../runtime/test_compositor_factory.py | 36 +- .../dify_agent/runtime/test_run_scheduler.py | 41 +- .../local/dify_agent/runtime/test_runner.py | 208 +- .../dify_agent/runtime_backend/test_e2b.py | 219 ++ .../test_enterprise_backend.py | 310 +++ .../dify_agent/runtime_backend/test_leases.py | 33 + .../dify_agent/runtime_backend/test_local.py | 446 ++++ .../runtime_backend/test_profile.py | 34 + .../runtime_backend/test_shellctl_backend.py | 191 ++ .../tests/local/dify_agent/server/test_app.py | 14 +- .../server/test_execution_bindings.py | 66 + .../dify_agent/server/test_home_snapshots.py | 116 ++ .../dify_agent/server/test_sandbox_files.py | 628 ------ .../local/dify_agent/server/test_settings.py | 105 +- .../dify_agent/server/test_workspace_files.py | 123 ++ dify-agent/tests/local/test_packaging.py | 1 + dify-agent/uv.lock | 54 + docker/.env.example | 14 +- docker/docker-compose-template.yaml | 15 +- docker/docker-compose.e2b.yaml | 50 + docker/docker-compose.yaml | 15 +- .../envs/core-services/dify-agent.env.example | 15 +- .../generated/api/console/agent/orpc.gen.ts | 8 +- .../generated/api/console/agent/types.gen.ts | 23 +- .../generated/api/console/agent/zod.gen.ts | 86 +- .../generated/api/console/apps/orpc.gen.ts | 3 +- .../generated/api/console/apps/types.gen.ts | 8 +- .../generated/api/console/apps/zod.gen.ts | 6 +- .../agent-orchestrate-panel-content.spec.tsx | 20 +- .../agent-orchestrate-panel-content.tsx | 28 +- .../configure/__tests__/page.spec.tsx | 4 +- .../configure/components/composer-session.tsx | 26 +- .../working-directory-panel.spec.tsx | 78 +- .../components/preview/chat-history.ts | 9 - .../components/preview/chat-runtime.tsx | 12 +- .../hook/use-working-directory-panel.tsx | 77 +- .../preview/working-directory-panel.tsx | 199 +- .../use-agent-configure-build-draft.ts | 1 + 214 files changed, 15160 insertions(+), 8787 deletions(-) delete mode 100644 api/clients/agent_backend/session_cleanup.py create mode 100644 api/core/workflow/node_execution_process_data.py delete mode 100644 api/core/workflow/nodes/agent_v2/session_cleanup_layer.py create mode 100644 api/core/workflow/nodes/agent_v2/workspace_retirement_layer.py create mode 100644 api/migrations/versions/2026_07_21_2251-2f39536b3feb_add_agent_home_snapshot_ledger.py create mode 100644 api/migrations/versions/2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py create mode 100644 api/services/agent/home_snapshot_service.py create mode 100644 api/services/agent/retirement_service.py create mode 100644 api/services/agent/workspace_service.py delete mode 100644 api/tasks/agent_backend_session_cleanup_task.py create mode 100644 api/tasks/collect_agent_resources_task.py delete mode 100644 api/tests/unit_tests/clients/agent_backend/test_cleanup_composition_compositor_integration.py delete mode 100644 api/tests/unit_tests/clients/agent_backend/test_session_cleanup.py delete mode 100644 api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py create mode 100644 api/tests/unit_tests/core/workflow/nodes/agent_v2/test_workspace_retirement_layer.py create mode 100644 api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py create mode 100644 api/tests/unit_tests/migrations/test_agent_workspace_migration.py create mode 100644 api/tests/unit_tests/services/agent/test_home_snapshot_service.py create mode 100644 api/tests/unit_tests/services/agent/test_retirement_service.py create mode 100644 api/tests/unit_tests/services/agent/test_workspace_service.py delete mode 100644 api/tests/unit_tests/tasks/test_agent_backend_session_cleanup_task.py create mode 100644 api/tests/unit_tests/tasks/test_collect_agent_resources_task.py create mode 100644 api/tests/unit_tests/tasks/test_workflow_node_execution_identity.py create mode 100644 dify-agent/docs/dify-agent/concepts/runtime-resources/index.md delete mode 100644 dify-agent/src/dify_agent/adapters/shell/config.py delete mode 100644 dify-agent/src/dify_agent/adapters/shell/enterprise/__init__.py delete mode 100644 dify-agent/src/dify_agent/adapters/shell/enterprise/enterprise.py delete mode 100644 dify-agent/src/dify_agent/adapters/shell/factory.py create mode 100644 dify-agent/src/dify_agent/layers/runtime/__init__.py create mode 100644 dify-agent/src/dify_agent/layers/runtime/configs.py create mode 100644 dify-agent/src/dify_agent/layers/runtime/layer.py create mode 100644 dify-agent/src/dify_agent/protocol/execution_binding.py create mode 100644 dify-agent/src/dify_agent/protocol/home_snapshot.py delete mode 100644 dify-agent/src/dify_agent/protocol/sandbox.py create mode 100644 dify-agent/src/dify_agent/protocol/workspace.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/__init__.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/e2b.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/enterprise.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/errors.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/leases.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/local.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/profile.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/protocols.py create mode 100644 dify-agent/src/dify_agent/runtime_backend/shellctl.py create mode 100644 dify-agent/src/dify_agent/server/execution_bindings.py create mode 100644 dify-agent/src/dify_agent/server/home_snapshots.py create mode 100644 dify-agent/src/dify_agent/server/routes/execution_bindings.py create mode 100644 dify-agent/src/dify_agent/server/routes/home_snapshots.py delete mode 100644 dify-agent/src/dify_agent/server/routes/sandbox_files.py create mode 100644 dify-agent/src/dify_agent/server/routes/workspace_files.py delete mode 100644 dify-agent/src/dify_agent/server/sandbox_files.py create mode 100644 dify-agent/src/dify_agent/server/workspace_files.py create mode 100755 dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh create mode 100644 dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py delete mode 100644 dify-agent/tests/local/dify_agent/adapters/shell/test_enterprise.py create mode 100644 dify-agent/tests/local/dify_agent/layers/test_runtime_layer.py delete mode 100644 dify-agent/tests/local/dify_agent/protocol/test_sandbox_locator.py create mode 100644 dify-agent/tests/local/dify_agent/protocol/test_working_environment.py create mode 100644 dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py create mode 100644 dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py create mode 100644 dify-agent/tests/local/dify_agent/runtime_backend/test_leases.py create mode 100644 dify-agent/tests/local/dify_agent/runtime_backend/test_local.py create mode 100644 dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py create mode 100644 dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py create mode 100644 dify-agent/tests/local/dify_agent/server/test_execution_bindings.py create mode 100644 dify-agent/tests/local/dify_agent/server/test_home_snapshots.py delete mode 100644 dify-agent/tests/local/dify_agent/server/test_sandbox_files.py create mode 100644 dify-agent/tests/local/dify_agent/server/test_workspace_files.py create mode 100644 docker/docker-compose.e2b.yaml diff --git a/.github/workflows/web-e2e.yml b/.github/workflows/web-e2e.yml index b0a81b01609..bbd86fcccc2 100644 --- a/.github/workflows/web-e2e.yml +++ b/.github/workflows/web-e2e.yml @@ -63,6 +63,7 @@ jobs: E2E_ADMIN_PASSWORD: E2eAdmin12345 E2E_FORCE_WEB_BUILD: "1" E2E_INIT_PASSWORD: E2eInit12345 + E2E_START_AGENT_BACKEND: "1" run: vp run e2e:full - name: Preserve Chromium E2E report and logs diff --git a/api/clients/agent_backend/__init__.py b/api/clients/agent_backend/__init__.py index 67175c4795c..dbdc93cfbad 100644 --- a/api/clients/agent_backend/__init__.py +++ b/api/clients/agent_backend/__init__.py @@ -5,8 +5,6 @@ API adapters: request building from Dify product concepts, a thin client wrapper event adaptation for future workflow integration, and deterministic fakes. """ -from dify_agent.protocol import RuntimeLayerSpec, extract_runtime_layer_specs - from clients.agent_backend.client import AgentBackendRunClient, DifyAgentBackendRunClient from clients.agent_backend.errors import ( AgentBackendError, @@ -47,11 +45,6 @@ from clients.agent_backend.request_builder import ( AgentBackendWorkflowNodeRunInput, redact_for_agent_backend_log, ) -from clients.agent_backend.session_cleanup import ( - AgentBackendSessionCleanupPayload, - AgentBackendSessionCleanupResult, - cleanup_agent_backend_session, -) __all__ = [ "AGENT_SOUL_PROMPT_LAYER_ID", @@ -80,8 +73,6 @@ __all__ = [ "AgentBackendRunRequestBuilder", "AgentBackendRunStartedInternalEvent", "AgentBackendRunSucceededInternalEvent", - "AgentBackendSessionCleanupPayload", - "AgentBackendSessionCleanupResult", "AgentBackendStreamError", "AgentBackendStreamInternalEvent", "AgentBackendTransportError", @@ -90,9 +81,6 @@ __all__ = [ "DifyAgentBackendRunClient", "FakeAgentBackendRunClient", "FakeAgentBackendScenario", - "RuntimeLayerSpec", - "cleanup_agent_backend_session", "create_agent_backend_run_client", - "extract_runtime_layer_specs", "redact_for_agent_backend_log", ] diff --git a/api/clients/agent_backend/request_builder.py b/api/clients/agent_backend/request_builder.py index 1fd7b2cd990..b55957e163a 100644 --- a/api/clients/agent_backend/request_builder.py +++ b/api/clients/agent_backend/request_builder.py @@ -16,7 +16,6 @@ from collections.abc import Mapping from typing import ClassVar, Literal from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot from agenton.layers import ExitIntent from agenton_collections.layers.plain import PLAIN_PROMPT_LAYER_TYPE_ID, PromptLayerConfig from agenton_collections.layers.pydantic_ai import PYDANTIC_AI_HISTORY_LAYER_TYPE_ID @@ -37,6 +36,7 @@ from dify_agent.layers.execution_context import ( ) from dify_agent.layers.knowledge import DIFY_KNOWLEDGE_BASE_LAYER_TYPE_ID, DifyKnowledgeBaseLayerConfig from dify_agent.layers.output import DIFY_OUTPUT_LAYER_TYPE_ID, DifyOutputLayerConfig +from dify_agent.layers.runtime import DIFY_RUNTIME_LAYER_TYPE_ID, DifyRuntimeLayerConfig from dify_agent.layers.shell import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig from dify_agent.protocol import ( DIFY_AGENT_HISTORY_LAYER_ID, @@ -47,7 +47,6 @@ from dify_agent.protocol import ( LayerExitSignals, RunComposition, RunLayerSpec, - RuntimeLayerSpec, ) from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator @@ -56,6 +55,7 @@ WORKFLOW_NODE_JOB_PROMPT_LAYER_ID = "workflow_node_job_prompt" WORKFLOW_USER_PROMPT_LAYER_ID = "workflow_user_prompt" AGENT_APP_USER_PROMPT_LAYER_ID = "agent_app_user_prompt" DIFY_EXECUTION_CONTEXT_LAYER_ID = "execution_context" +DIFY_RUNTIME_LAYER_ID = "runtime" DIFY_CONFIG_LAYER_ID = "config" DIFY_DRIVE_LAYER_ID = "drive" DIFY_PLUGIN_TOOLS_LAYER_ID = "tools" @@ -66,25 +66,11 @@ DIFY_SHELL_LAYER_ID = "shell" type AgentConfigVersionKind = Literal["snapshot", "draft", "build_draft"] -def _filter_snapshot_to_specs( - snapshot: CompositorSessionSnapshot, - specs: list[RuntimeLayerSpec], -) -> CompositorSessionSnapshot: - """Keep only snapshot layers whose names appear in the cleanup spec list. - - The agenton compositor rejects a snapshot whose layer-name sequence does - not match the active composition exactly. Cleanup-replay drops plugin - layers, so we must drop the matching snapshot entries here. - """ - kept_names = {spec.name for spec in specs} - filtered_layers: list[LayerSessionSnapshot] = [layer for layer in snapshot.layers if layer.name in kept_names] - if len(filtered_layers) == len(snapshot.layers): - return snapshot - return CompositorSessionSnapshot(schema_version=snapshot.schema_version, layers=filtered_layers) - - def _shell_layer_deps() -> dict[str, str]: - return {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID} + return { + "execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID, + "runtime": DIFY_RUNTIME_LAYER_ID, + } def _drive_layer_deps() -> dict[str, str]: @@ -214,6 +200,7 @@ class AgentBackendWorkflowNodeRunInput(BaseModel): model: AgentBackendModelConfig execution_context: DifyExecutionContextLayerConfig + backend_binding_ref: str = Field(min_length=1) workflow_node_job_prompt: str user_prompt: str agent_soul_prompt: str | None = None @@ -231,8 +218,8 @@ class AgentBackendWorkflowNodeRunInput(BaseModel): # the Agent Soul configures human involvement; a deferred call ends the run and # the workflow pauses via the existing HITL form mechanism (ENG-635). ask_human_config: DifyAskHumanLayerConfig | None = None - # Inject the sandboxed shell layer (dify.shell). Requires the agent backend - # to be wired with a shellctl entrypoint; see configs AGENT_SHELL_ENABLED. + # Inject the sandboxed shell graph. Requires a deployment-selected runtime + # backend plus the product-resolved persistent Binding. include_shell: bool = False shell_config: DifyShellLayerConfig | None = None session_snapshot: CompositorSessionSnapshot | None = None @@ -240,7 +227,6 @@ class AgentBackendWorkflowNodeRunInput(BaseModel): # (ENG-638). Keyed by the original deferred tool_call_id. deferred_tool_results: DeferredToolResultsPayload | None = None include_history: bool = True - suspend_on_exit: bool = True metadata: dict[str, JsonValue] = Field(default_factory=dict) model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid", arbitrary_types_allowed=True) @@ -264,6 +250,7 @@ class AgentBackendAgentAppRunInput(BaseModel): model: AgentBackendModelConfig execution_context: DifyExecutionContextLayerConfig + backend_binding_ref: str = Field(min_length=1) user_prompt: str agent_soul_prompt: str | None = None agent_config_version_kind: AgentConfigVersionKind = "snapshot" @@ -279,8 +266,8 @@ class AgentBackendAgentAppRunInput(BaseModel): # Human-in-the-loop ask_human deferred tool (dify.ask_human). Present only when # the Agent Soul configures human involvement (ENG-635). ask_human_config: DifyAskHumanLayerConfig | None = None - # Inject the sandboxed shell layer (dify.shell). Requires the agent backend - # to be wired with a shellctl entrypoint; see configs AGENT_SHELL_ENABLED. + # Inject the sandboxed shell graph. Requires a deployment-selected runtime + # backend plus the product-resolved persistent Binding. include_shell: bool = False shell_config: DifyShellLayerConfig | None = None session_snapshot: CompositorSessionSnapshot | None = None @@ -288,7 +275,6 @@ class AgentBackendAgentAppRunInput(BaseModel): # (ENG-638). Keyed by the original deferred tool_call_id. deferred_tool_results: DeferredToolResultsPayload | None = None include_history: bool = True - suspend_on_exit: bool = True metadata: dict[str, JsonValue] = Field(default_factory=dict) model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid", arbitrary_types_allowed=True) @@ -350,6 +336,14 @@ class AgentBackendRunRequestBuilder: run_input.include_shell or run_input.config_layer_config is not None or run_input.drive_config is not None ) if include_shell: + layers.append( + RunLayerSpec( + name=DIFY_RUNTIME_LAYER_ID, + type=DIFY_RUNTIME_LAYER_TYPE_ID, + metadata=run_input.metadata, + config=DifyRuntimeLayerConfig(backend_binding_ref=run_input.backend_binding_ref), + ) + ) # Sandboxed bash workspace (dify.shell). It enters before config/drive # so eager pulls materialize content in the same filesystem used by # model commands. @@ -481,53 +475,7 @@ class AgentBackendRunRequestBuilder: metadata=run_input.metadata, session_snapshot=run_input.session_snapshot, deferred_tool_results=run_input.deferred_tool_results, - on_exit=LayerExitSignals( - default=ExitIntent.SUSPEND if run_input.suspend_on_exit else ExitIntent.DELETE, - ), - ) - - def build_cleanup_request( - self, - *, - session_snapshot: CompositorSessionSnapshot, - runtime_layer_specs: list[RuntimeLayerSpec], - idempotency_key: str | None = None, - metadata: dict[str, JsonValue] | None = None, - ) -> CreateRunRequest: - """Build a lifecycle-only cleanup request that replays the prior layers. - - The agenton compositor enforces that the session snapshot's layer names - match the active composition in order, so cleanup must replay the same - non-plugin layer graph that produced the snapshot. Plugin layers - (``dify.plugin.llm``, ``dify.plugin.tools``) are excluded from both the - composition and the snapshot before submission because their configs - may carry credentials or runtime-only declarations that are not - persisted between runs. - """ - if not runtime_layer_specs: - raise ValueError( - "build_cleanup_request requires runtime_layer_specs; an empty " - "composition would fail the agent backend's snapshot validation." - ) - request_metadata = dict(metadata or {}) - request_metadata["agent_backend_lifecycle"] = "session_cleanup" - layers = [ - RunLayerSpec( - name=spec.name, - type=spec.type, - deps=dict(spec.deps), - metadata=dict(spec.metadata), - config=spec.config, - ) - for spec in runtime_layer_specs - ] - filtered_snapshot = _filter_snapshot_to_specs(session_snapshot, runtime_layer_specs) - return CreateRunRequest( - composition=RunComposition(layers=layers), - idempotency_key=idempotency_key, - metadata=request_metadata, - session_snapshot=filtered_snapshot, - on_exit=LayerExitSignals(default=ExitIntent.DELETE), + on_exit=LayerExitSignals(default=ExitIntent.SUSPEND), ) def build_for_workflow_node(self, run_input: AgentBackendWorkflowNodeRunInput) -> CreateRunRequest: @@ -580,6 +528,14 @@ class AgentBackendRunRequestBuilder: run_input.include_shell or run_input.config_layer_config is not None or run_input.drive_config is not None ) if include_shell: + layers.append( + RunLayerSpec( + name=DIFY_RUNTIME_LAYER_ID, + type=DIFY_RUNTIME_LAYER_TYPE_ID, + metadata=run_input.metadata, + config=DifyRuntimeLayerConfig(backend_binding_ref=run_input.backend_binding_ref), + ) + ) # Sandboxed bash workspace (dify.shell). It enters before drive so # drive can materialize mentioned targets with `dify-agent drive pull` # in the same shell-visible filesystem used by model commands. @@ -713,9 +669,7 @@ class AgentBackendRunRequestBuilder: metadata=run_input.metadata, session_snapshot=run_input.session_snapshot, deferred_tool_results=run_input.deferred_tool_results, - on_exit=LayerExitSignals( - default=ExitIntent.SUSPEND if run_input.suspend_on_exit else ExitIntent.DELETE, - ), + on_exit=LayerExitSignals(default=ExitIntent.SUSPEND), ) diff --git a/api/clients/agent_backend/session_cleanup.py b/api/clients/agent_backend/session_cleanup.py deleted file mode 100644 index 370ff544e68..00000000000 --- a/api/clients/agent_backend/session_cleanup.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Shared API-side helper for Agent backend lifecycle-only session cleanup. - -Product code owns local row retirement and background-task dispatch. This module -only adapts persisted cleanup inputs into the public ``dify-agent`` run -protocol, performs the synchronous ``create_run + wait_run`` loop used by Celery -workers, and reports whether the backend cleanup succeeded, was skipped, or -failed. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import ClassVar, Literal - -from agenton.compositor import CompositorSessionSnapshot -from dify_agent.protocol import RuntimeLayerSpec -from pydantic import BaseModel, ConfigDict, Field, JsonValue - -from clients.agent_backend.client import AgentBackendRunClient -from clients.agent_backend.errors import AgentBackendError -from clients.agent_backend.request_builder import AgentBackendRunRequestBuilder - - -class AgentBackendSessionCleanupPayload(BaseModel): - """Serialized cleanup inputs preserved across API and Celery boundaries.""" - - session_snapshot: CompositorSessionSnapshot | None = None - runtime_layer_specs: list[RuntimeLayerSpec] = Field(default_factory=list) - idempotency_key: str | None = None - metadata: dict[str, JsonValue] = Field(default_factory=dict) - timeout_seconds: float = 30.0 - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -@dataclass(frozen=True, slots=True) -class AgentBackendSessionCleanupResult: - """Terminal outcome of one backend cleanup attempt.""" - - status: Literal["succeeded", "skipped", "failed"] - reason: str | None = None - cleanup_run_id: str | None = None - - @classmethod - def succeeded(cls, cleanup_run_id: str) -> AgentBackendSessionCleanupResult: - return cls(status="succeeded", cleanup_run_id=cleanup_run_id) - - @classmethod - def skipped(cls, reason: str) -> AgentBackendSessionCleanupResult: - return cls(status="skipped", reason=reason) - - @classmethod - def failed(cls, reason: str, cleanup_run_id: str | None = None) -> AgentBackendSessionCleanupResult: - return cls(status="failed", reason=reason, cleanup_run_id=cleanup_run_id) - - -def cleanup_agent_backend_session( - *, - payload: AgentBackendSessionCleanupPayload, - client: AgentBackendRunClient | None, - request_builder: AgentBackendRunRequestBuilder | None = None, -) -> AgentBackendSessionCleanupResult: - """Run lifecycle-only cleanup against the Agent backend and report status.""" - if client is None: - return AgentBackendSessionCleanupResult.skipped("no_agent_backend_client") - if payload.session_snapshot is None: - return AgentBackendSessionCleanupResult.skipped("missing_session_snapshot") - if not payload.runtime_layer_specs: - return AgentBackendSessionCleanupResult.skipped("missing_runtime_layer_specs") - - builder = request_builder or AgentBackendRunRequestBuilder() - request = builder.build_cleanup_request( - session_snapshot=payload.session_snapshot, - runtime_layer_specs=payload.runtime_layer_specs, - idempotency_key=payload.idempotency_key, - metadata=payload.metadata, - ) - - try: - response = client.create_run(request) - except AgentBackendError as exc: - return AgentBackendSessionCleanupResult.failed(str(exc)) - - try: - status_response = client.wait_run(response.run_id, timeout_seconds=payload.timeout_seconds) - except AgentBackendError as exc: - return AgentBackendSessionCleanupResult.failed(str(exc), cleanup_run_id=response.run_id) - - if status_response.status != "succeeded": - reason = status_response.error or f"cleanup run ended with status {status_response.status}" - return AgentBackendSessionCleanupResult.failed(reason, cleanup_run_id=response.run_id) - - return AgentBackendSessionCleanupResult.succeeded(response.run_id) - - -__all__ = [ - "AgentBackendSessionCleanupPayload", - "AgentBackendSessionCleanupResult", - "cleanup_agent_backend_session", -] diff --git a/api/configs/extra/agent_backend_config.py b/api/configs/extra/agent_backend_config.py index d5caf3d2e3c..fd412936390 100644 --- a/api/configs/extra/agent_backend_config.py +++ b/api/configs/extra/agent_backend_config.py @@ -44,9 +44,8 @@ class AgentBackendConfig(BaseSettings): AGENT_SHELL_ENABLED: bool = Field( description=( - "Inject the dify.shell layer (sandboxed bash workspace) into Agent runs. " - "Requires the agent backend to be wired with a shellctl entrypoint before " - "shell-using Agent runs are executed." + "Inject the Home, Workspace, Sandbox, and Shell runtime layers into Agent runs. " + "Requires Dify Agent to have a deployment-selected runtime backend." ), default=True, ) diff --git a/api/controllers/console/agent/roster.py b/api/controllers/console/agent/roster.py index ea313ff37aa..2892e30cb05 100644 --- a/api/controllers/console/agent/roster.py +++ b/api/controllers/console/agent/roster.py @@ -62,7 +62,7 @@ from libs.datetime_utils import parse_time_range from libs.helper import dump_response from libs.login import login_required from models import Account -from models.agent import Agent, AgentConfigDraftType, AgentStatus +from models.agent import Agent, AgentStatus from models.agent_config_entities import AgentSoulConfig from models.enums import ApiTokenType from models.model import ApiToken, App, IconType @@ -265,13 +265,6 @@ class AgentDebugConversationRefreshResponse(BaseModel): debug_conversation_message_count: int = 0 -class AgentDebugConversationRefreshPayload(BaseModel): - draft_type: AgentConfigDraftType = Field( - default=AgentConfigDraftType.DEBUG_BUILD, - description="Agent draft surface whose conversation should be refreshed", - ) - - class AgentPublishPayload(BaseModel): version_note: str | None = Field(default=None, description="Optional note for this published Agent version") @@ -315,7 +308,6 @@ register_schema_models( AgentAppCopyPayload, AgentPublishPayload, AgentBuildDraftCheckoutPayload, - AgentDebugConversationRefreshPayload, ComposerSavePayload, AgentApiStatusPayload, AgentInviteOptionsQuery, @@ -395,11 +387,10 @@ def _serialize_agent_app_detail( payload["backing_app_id"] = roster_service.runtime_backing_app_id(agent) payload["hidden_app_backed"] = bool(agent.backing_app_id and agent.backing_app_id != agent.app_id) payload["id"] = agent.id - debug_conversation_id = roster_service.get_or_create_agent_app_debug_conversation_id( + debug_conversation_id = roster_service.get_or_create_build_conversation( tenant_id=app_model.tenant_id, agent_id=agent.id, account_id=current_user.id, - draft_type=AgentConfigDraftType.DEBUG_BUILD, commit=False, ) message_count = roster_service.count_agent_app_debug_conversation_messages( @@ -439,11 +430,10 @@ def _serialize_agent_app_pagination(session: Session, app_pagination, *, tenant_ tenant_id=tenant_id, agent_ids=[agent.id for agent in agents_by_app_id.values()], ) - debug_conversation_ids_by_agent_id = roster_service.load_or_create_agent_app_debug_conversation_ids_by_agent_id( + debug_conversation_ids_by_agent_id = roster_service.load_or_create_build_conversation_ids_by_agent_id( tenant_id=tenant_id, agents=list(agents_by_app_id.values()), account_id=current_user.id, - draft_type=AgentConfigDraftType.DEBUG_BUILD, ) payload = AgentAppPagination.model_validate( app_pagination, @@ -678,16 +668,6 @@ class AgentAppApi(Resource): @console_ns.route("/agent//debug-conversation/refresh") class AgentDebugConversationRefreshApi(Resource): - @console_ns.expect(console_ns.models[AgentDebugConversationRefreshPayload.__name__]) - @console_ns.doc( - params={ - "payload": { - "in": "body", - "required": False, - "schema": {"$ref": f"#/components/schemas/{AgentDebugConversationRefreshPayload.__name__}"}, - } - } - ) @console_ns.response( 200, "Agent debug conversation refreshed", @@ -702,12 +682,10 @@ class AgentDebugConversationRefreshApi(Resource): @with_current_tenant_id @with_session def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID): - args = AgentDebugConversationRefreshPayload.model_validate(request.get_json(silent=True) or {}) - debug_conversation_id = _agent_roster_service(session).refresh_agent_app_debug_conversation_id( + debug_conversation_id = _agent_roster_service(session).reset_build_conversation( tenant_id=tenant_id, agent_id=str(agent_id), account_id=current_user.id, - draft_type=args.draft_type, ) return AgentDebugConversationRefreshResponse( debug_conversation_id=debug_conversation_id, @@ -751,7 +729,7 @@ class AgentBuildDraftCheckoutApi(Resource): @rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.AGENT_MANAGE, resource_required=False) @with_current_user @with_current_tenant_id - @with_session + @with_session(write=False) def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID): args = AgentBuildDraftCheckoutPayload.model_validate(console_ns.payload or {}) return AgentComposerService.checkout_agent_app_build_draft( @@ -808,7 +786,7 @@ class AgentBuildDraftApi(Resource): @edit_permission_required @with_current_user @with_current_tenant_id - @with_session + @with_session(write=False) def delete(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID): return AgentComposerService.discard_agent_app_build_draft( session=session, @@ -828,7 +806,7 @@ class AgentBuildDraftApplyApi(Resource): @rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.AGENT_MANAGE, resource_required=False) @with_current_user @with_current_tenant_id - @with_session + @with_session(write=False) def post(self, session: Session, tenant_id: str, current_user: Account, agent_id: UUID): return AgentComposerService.apply_agent_app_build_draft( session=session, diff --git a/api/controllers/console/app/agent_app_sandbox.py b/api/controllers/console/app/agent_app_sandbox.py index 17b5bcd23ad..6c37543316a 100644 --- a/api/controllers/console/app/agent_app_sandbox.py +++ b/api/controllers/console/app/agent_app_sandbox.py @@ -1,7 +1,7 @@ """Console routes for Agent App and workflow Agent sandbox file access. -The API keeps product-facing locators (conversation or workflow node identity) -on this public boundary and proxies list/read/upload to the agent backend's new +The API accepts product-facing Conversation, Build Draft, or Workflow Node +Execution locators and proxies list/read/upload to the agent backend's ``/sandbox`` contract. """ @@ -26,10 +26,16 @@ from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.agent.app_helpers import resolve_agent_runtime_app_model from controllers.console.app.wraps import get_app_model -from controllers.console.wraps import account_initialization_required, setup_required, with_current_tenant_id +from controllers.console.wraps import ( + account_initialization_required, + setup_required, + with_current_tenant_id, + with_current_user, +) from extensions.ext_database import db from fields.base import ResponseModel from libs.login import login_required +from models import Account from models.model import App, AppMode from services.agent_app_sandbox_service import ( AgentAppSandboxService, @@ -37,52 +43,43 @@ from services.agent_app_sandbox_service import ( WorkflowAgentSandboxService, ) -_NODE_EXECUTION_ID_DESCRIPTION = ( - "Optional workflow node execution ID. When omitted, the latest active session for the node is used." -) - class AgentSandboxListQuery(BaseModel): - conversation_id: str = Field(min_length=1, description="Agent App conversation ID") + caller_type: Literal["conversation", "build_draft"] + caller_id: str = Field(min_length=1, description="Agent App caller ID") path: str = Field(default=".", description="Directory path relative to the sandbox workspace") class AgentSandboxInfoQuery(BaseModel): - conversation_id: str = Field(min_length=1, description="Agent App conversation ID") + caller_type: Literal["conversation", "build_draft"] + caller_id: str = Field(min_length=1, description="Agent App caller ID") class AgentSandboxFileQuery(BaseModel): - conversation_id: str = Field(min_length=1, description="Agent App conversation ID") + caller_type: Literal["conversation", "build_draft"] + caller_id: str = Field(min_length=1, description="Agent App caller ID") path: str = Field(min_length=1, description="File path relative to the sandbox workspace") class AgentSandboxUploadPayload(BaseModel): - conversation_id: str = Field(min_length=1, description="Agent App conversation ID") + caller_type: Literal["conversation", "build_draft"] + caller_id: str = Field(min_length=1, description="Agent App caller ID") path: str = Field(min_length=1, description="File path relative to the sandbox workspace") class WorkflowAgentSandboxListQuery(BaseModel): + node_execution_id: str = Field(min_length=1, description="Workflow node execution ID") path: str = Field(default=".", description="Directory path relative to the sandbox workspace") - node_execution_id: str | None = Field( - default=None, - description=_NODE_EXECUTION_ID_DESCRIPTION, - ) class WorkflowAgentSandboxFileQuery(BaseModel): + node_execution_id: str = Field(min_length=1, description="Workflow node execution ID") path: str = Field(min_length=1, description="File path relative to the sandbox workspace") - node_execution_id: str | None = Field( - default=None, - description=_NODE_EXECUTION_ID_DESCRIPTION, - ) class WorkflowAgentSandboxUploadPayload(BaseModel): + node_execution_id: str = Field(min_length=1, description="Workflow node execution ID") path: str = Field(min_length=1, description="File path relative to the sandbox workspace") - node_execution_id: str | None = Field( - default=None, - description=_NODE_EXECUTION_ID_DESCRIPTION, - ) class SandboxFileEntryResponse(ResponseModel): @@ -99,7 +96,6 @@ class SandboxListResponse(ResponseModel): class SandboxInfoResponse(ResponseModel): - session_id: str workspace_cwd: str @@ -155,15 +151,19 @@ class AgentAppSandboxInfoResource(Resource): @login_required @account_initialization_required @with_current_tenant_id + @with_current_user @with_session(write=False) - def get(self, session: Session, tenant_id: str, agent_id: UUID): + def get(self, session: Session, current_user: Account, tenant_id: str, agent_id: UUID): app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id) query = query_params_from_request(AgentSandboxInfoQuery) try: result = AgentAppSandboxService().get_info( tenant_id=tenant_id, app_id=app_model.id, - conversation_id=query.conversation_id, + agent_id=str(agent_id), + caller_type=query.caller_type, + caller_id=query.caller_id, + account_id=current_user.id, ) except Exception as exc: return _handle(exc) @@ -180,15 +180,19 @@ class AgentAppSandboxListResource(Resource): @login_required @account_initialization_required @with_current_tenant_id + @with_current_user @with_session(write=False) - def get(self, session: Session, tenant_id: str, agent_id: UUID): + def get(self, session: Session, current_user: Account, tenant_id: str, agent_id: UUID): app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id) query = query_params_from_request(AgentSandboxListQuery) try: result = AgentAppSandboxService().list_files( tenant_id=tenant_id, app_id=app_model.id, - conversation_id=query.conversation_id, + agent_id=str(agent_id), + caller_type=query.caller_type, + caller_id=query.caller_id, + account_id=current_user.id, path=query.path, ) except Exception as exc: @@ -206,15 +210,19 @@ class AgentAppSandboxReadResource(Resource): @login_required @account_initialization_required @with_current_tenant_id + @with_current_user @with_session(write=False) - def get(self, session: Session, tenant_id: str, agent_id: UUID): + def get(self, session: Session, current_user: Account, tenant_id: str, agent_id: UUID): app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id) query = query_params_from_request(AgentSandboxFileQuery) try: result = AgentAppSandboxService().read_file( tenant_id=tenant_id, app_id=app_model.id, - conversation_id=query.conversation_id, + agent_id=str(agent_id), + caller_type=query.caller_type, + caller_id=query.caller_id, + account_id=current_user.id, path=query.path, ) except Exception as exc: @@ -232,15 +240,19 @@ class AgentAppSandboxUploadResource(Resource): @login_required @account_initialization_required @with_current_tenant_id + @with_current_user @with_session(write=False) - def post(self, session: Session, tenant_id: str, agent_id: UUID): + def post(self, session: Session, current_user: Account, tenant_id: str, agent_id: UUID): app_model = resolve_agent_runtime_app_model(session=session, tenant_id=tenant_id, agent_id=agent_id) payload = AgentSandboxUploadPayload.model_validate(request.get_json(silent=True) or {}) try: result = AgentAppSandboxService().upload_file( tenant_id=tenant_id, app_id=app_model.id, - conversation_id=payload.conversation_id, + agent_id=str(agent_id), + caller_type=payload.caller_type, + caller_id=payload.caller_id, + account_id=current_user.id, path=payload.path, ) except Exception as exc: diff --git a/api/controllers/console/app/completion.py b/api/controllers/console/app/completion.py index beca00d621b..8fc8d12ce3a 100644 --- a/api/controllers/console/app/completion.py +++ b/api/controllers/console/app/completion.py @@ -36,7 +36,7 @@ from controllers.console.wraps import ( with_current_user_id, ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError -from core.app.entities.app_invoke_entities import AGENT_RUNTIME_EXIT_INTENT_ARG, InvokeFrom +from core.app.entities.app_invoke_entities import InvokeFrom from core.app.features.rate_limiting.rate_limit import RateLimitGenerator from core.errors.error import ( ModelCurrentlyNotSupportError, @@ -353,12 +353,7 @@ def _resolve_current_user_agent_debug_conversation_id( draft_type: AgentConfigDraftType, start_new: bool = False, ) -> str: - """Resolve or rotate the current editor's conversation within one draft surface. - - ``start_new`` rotates the scoped mapping through ``AgentRosterService`` so - the old runtime session is retired before the new conversation is used. - Continuations and Build chat keep resolving the existing mapping. - """ + """Resolve the current editor's Build or Preview conversation.""" roster_service = AgentRosterService(session) resolved_agent_id = agent_id @@ -368,17 +363,26 @@ def _resolve_current_user_agent_debug_conversation_id( raise AgentNotFoundError() resolved_agent_id = agent.id - resolve_conversation = ( - roster_service.refresh_agent_app_debug_conversation_id - if start_new - else roster_service.get_or_create_agent_app_debug_conversation_id - ) - return resolve_conversation( + if draft_type == AgentConfigDraftType.DEBUG_BUILD: + return roster_service.get_or_create_build_conversation( + tenant_id=current_tenant_id, + agent_id=resolved_agent_id, + account_id=current_user.id, + ) + if start_new: + return roster_service.rotate_preview_conversation( + tenant_id=current_tenant_id, + agent_id=resolved_agent_id, + account_id=current_user.id, + ) + conversation_id = roster_service.get_current_preview_conversation( tenant_id=current_tenant_id, agent_id=resolved_agent_id, account_id=current_user.id, - draft_type=draft_type, ) + if conversation_id is None: + raise NotFound("Conversation Not Exists.") + return conversation_id def _create_chat_message( @@ -450,7 +454,6 @@ def _create_build_chat_finalization_message( "draft_type": "debug_build", "conversation_id": debug_conversation_id, "auto_generate_name": False, - AGENT_RUNTIME_EXIT_INTENT_ARG: "delete", } external_trace_id = get_external_trace_id(request) if external_trace_id: diff --git a/api/controllers/console/app/workflow.py b/api/controllers/console/app/workflow.py index f934f440b0e..958bf830c87 100644 --- a/api/controllers/console/app/workflow.py +++ b/api/controllers/console/app/workflow.py @@ -76,11 +76,13 @@ from models import Account, App from models.model import AppMode from models.workflow import Workflow from repositories.workflow_collaboration_repository import WORKFLOW_ONLINE_USERS_PREFIX +from services.agent.retirement_service import WorkflowAgentRetirementService from services.app_generate_service import AppGenerateService from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.errors.llm import InvokeRateLimitError from services.workflow_ref_service import WorkflowRefService from services.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError, WorkflowService +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -1245,7 +1247,7 @@ class PublishedWorkflowApi(Resource): workflow_service = WorkflowService() with sessionmaker(db.engine).begin() as session: - workflow = workflow_service.publish_workflow( + workflow, retirement_candidates = workflow_service.publish_workflow( session=session, app_model=app_model, account=current_user, @@ -1262,6 +1264,16 @@ class PublishedWorkflowApi(Resource): workflow_created_at = TimestampField().format(workflow.created_at) + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=app_model.tenant_id, + agent_ids=retirement_candidates, + account_id=current_user.id, + ) + enqueue_agent_resource_collection( + tenant_id=app_model.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) return { "result": "success", "created_at": workflow_created_at, diff --git a/api/controllers/console/snippets/snippet_workflow.py b/api/controllers/console/snippets/snippet_workflow.py index 558d50930c5..295170d2c04 100644 --- a/api/controllers/console/snippets/snippet_workflow.py +++ b/api/controllers/console/snippets/snippet_workflow.py @@ -56,10 +56,12 @@ from libs.helper import TimestampField from libs.login import current_account_with_tenant, login_required from models import Account from models.snippet import CustomizedSnippet +from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.workflow_publish_service import WorkflowAgentPublishService from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.snippet_generate_service import SnippetGenerateService from services.snippet_service import SnippetService +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -295,8 +297,9 @@ class SnippetPublishedWorkflowApi(Resource): with Session(db.engine) as session: snippet = session.merge(snippet) + tenant_id = snippet.tenant_id try: - workflow = snippet_service.publish_workflow( + workflow, retirement_candidates = snippet_service.publish_workflow( session=session, snippet=snippet, account=current_user, @@ -306,6 +309,16 @@ class SnippetPublishedWorkflowApi(Resource): except ValueError as e: return {"message": str(e)}, 400 + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=tenant_id, + agent_ids=retirement_candidates, + account_id=current_user.id, + ) + enqueue_agent_resource_collection( + tenant_id=tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) return { "result": "success", "created_at": workflow_created_at, diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py index 31a65578d49..cf3943a6ccc 100644 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ b/api/core/app/apps/advanced_chat/app_runner.py @@ -16,6 +16,7 @@ from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner from core.app.entities.app_invoke_entities import ( AdvancedChatAppGenerateEntity, AppGenerateEntity, + DifyRunContext, InvokeFrom, ) from core.app.entities.queue_entities import ( @@ -31,7 +32,7 @@ from core.moderation.base import ModerationError from core.moderation.input_moderation import InputModeration from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.workflow.node_factory import get_default_root_node_id -from core.workflow.nodes.agent_v2.session_cleanup_layer import build_workflow_agent_session_cleanup_layer +from core.workflow.nodes.agent_v2.workspace_retirement_layer import build_workflow_agent_workspace_retirement_layer from core.workflow.system_variables import ( build_bootstrap_variables, build_system_variables, @@ -267,7 +268,18 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner): ) workflow_entry.graph_engine.layer(persistence_layer) - workflow_entry.graph_engine.layer(build_workflow_agent_session_cleanup_layer()) + workflow_entry.graph_engine.layer( + build_workflow_agent_workspace_retirement_layer( + dify_run_context=DifyRunContext( + tenant_id=self._workflow.tenant_id, + app_id=self._workflow.app_id, + user_id=self.application_generate_entity.user_id, + user_from=user_from, + invoke_from=invoke_from, + trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + ) + ) + ) conversation_variable_layer = ConversationVariablePersistenceLayer( ConversationVariableUpdater(session_factory.get_session_maker()) ) diff --git a/api/core/app/apps/agent_app/app_generator.py b/api/core/app/apps/agent_app/app_generator.py index 2f343202f4b..a4337db2a8f 100644 --- a/api/core/app/apps/agent_app/app_generator.py +++ b/api/core/app/apps/agent_app/app_generator.py @@ -33,15 +33,13 @@ from core.app.apps.agent_app.app_runner import AgentAppRunner from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError from core.app.apps.agent_app.generate_response_converter import AgentAppGenerateResponseConverter from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeRequestBuilder -from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore +from core.app.apps.agent_app.session_store import AgentAppWorkspaceStore from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.message_based_app_generator import MessageBasedAppGenerator from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import ( - AGENT_RUNTIME_EXIT_INTENT_ARG, AgentAppGenerateEntity, - AgentRuntimeExitIntent, DifyRunContext, InvokeFrom, UserFrom, @@ -51,19 +49,24 @@ from core.db.session_factory import session_factory from core.ops.ops_trace_manager import TraceQueueManager from core.workflow.file_reference import build_file_reference, is_canonical_file_reference from extensions.ext_database import db -from models import Account, App, AppModelConfig, EndUser, Message, MessageAnnotation +from models import Account, App, AppModelConfig, Conversation, EndUser, Message, MessageAnnotation from models.agent import ( APP_BACKED_AGENT_SOURCES, Agent, AgentConfigDraft, AgentConfigDraftType, AgentConfigSnapshot, + AgentConfigVersionKind, AgentScope, AgentSource, AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, ) from models.agent_config_entities import AgentSoulConfig from models.model import load_annotation_reply_config +from services.agent.workspace_service import AgentWorkspaceService, WorkspaceOwnerScope from services.conversation_service import ConversationService logger = logging.getLogger(__name__) @@ -150,19 +153,6 @@ class AgentAppGenerator(MessageBasedAppGenerator): inputs = args["inputs"] prompt_file_mappings = args.get("files") or [] - # Resolve the bound roster Agent + its current Agent Soul snapshot. - agent, agent_config_id, agent_config_version_kind, agent_soul = self._resolve_agent( - app_model, - invoke_from=invoke_from, - draft_type=args.get("draft_type"), - user=user, - session=session, - ) - runtime_session_snapshot_id = self._runtime_session_snapshot_id( - invoke_from=invoke_from, - snapshot_id=agent_config_id, - ) - conversation = None conversation_id = args.get("conversation_id") if conversation_id: @@ -170,6 +160,21 @@ class AgentAppGenerator(MessageBasedAppGenerator): app_model=app_model, conversation_id=conversation_id, user=user, session=session ) + # New conversations use the current Agent generation. Existing + # conversations use the immutable generation named by their Binding. + agent, agent_config_id, agent_config_version_kind, agent_soul = self._resolve_agent( + app_model, + invoke_from=invoke_from, + draft_type=args.get("draft_type"), + user=user, + session=session, + conversation=conversation, + ) + session_scope_config_version_id = self._session_scope_config_version_id( + invoke_from=invoke_from, + config_version_id=agent_config_id, + ) + # Build the EasyUI-shaped config from the Agent Soul so the chat pipeline # can persist usage; the answer itself comes from the agent backend. app_model_config = ( @@ -186,8 +191,6 @@ class AgentAppGenerator(MessageBasedAppGenerator): model_conf = ModelConfigConverter.convert(app_config) trace_manager = TraceQueueManager(app_model.id, user.id if isinstance(user, Account) else user.session_id) - agent_runtime_exit_intent = self._resolve_agent_runtime_exit_intent(args) - application_generate_entity = AgentAppGenerateEntity( task_id=str(uuid.uuid4()), app_config=app_config, @@ -215,8 +218,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): agent_id=agent.id, agent_config_snapshot_id=agent_config_id, agent_config_version_kind=agent_config_version_kind, - agent_runtime_session_snapshot_id=runtime_session_snapshot_id, - agent_runtime_exit_intent=agent_runtime_exit_intent, + agent_session_scope_config_version_id=session_scope_config_version_id, ) conversation, message = self._init_generate_records( @@ -265,6 +267,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): app_model: App, user: Account | EndUser, conversation_id: str, + form_id: str, invoke_from: InvokeFrom, session: Session, ) -> None: @@ -279,14 +282,21 @@ class AgentAppGenerator(MessageBasedAppGenerator): conversation = ConversationService.get_conversation( app_model=app_model, conversation_id=conversation_id, user=user, session=session ) + draft_type, draft_id = self._resolve_resume_draft( + app_model=app_model, + conversation=conversation, + user=user, + form_id=form_id, + session=session, + ) agent, agent_config_id, agent_config_version_kind, agent_soul = self._resolve_agent( app_model, invoke_from=invoke_from, - draft_type=self._resume_draft_type( - app_model=app_model, conversation=conversation, user=user, session=session - ), + draft_type=draft_type, + draft_id=draft_id, user=user, session=session, + conversation=conversation, ) app_model_config = ( @@ -384,30 +394,37 @@ class AgentAppGenerator(MessageBasedAppGenerator): ) @staticmethod - def _resume_draft_type( - *, app_model: App, conversation: Any, user: Account | EndUser, session: Session - ) -> str | None: + def _resolve_resume_draft( + *, + app_model: App, + conversation: Any, + user: Account | EndUser, + form_id: str, + session: Session, + ) -> tuple[str | None, str | None]: if conversation.invoke_from != InvokeFrom.DEBUGGER: - return None - active_session = AgentAppRuntimeSessionStore().load_active_session_for_conversation( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - conversation_id=conversation.id, - ) - snapshot_id = active_session.scope.agent_config_snapshot_id if active_session is not None else None - if snapshot_id and isinstance(user, Account): - draft = session.scalar( - select(AgentConfigDraft).where( - AgentConfigDraft.tenant_id == app_model.tenant_id, - AgentConfigDraft.id == snapshot_id, - ) + return None, None + if not isinstance(user, Account): + return AgentConfigDraftType.DRAFT.value, None + + build_draft = session.scalar( + select(AgentConfigDraft) + .join( + AgentWorkspaceBinding, + AgentWorkspaceBinding.id == AgentConfigDraft.agent_workspace_binding_id, ) - if draft is not None: - if draft.draft_type == AgentConfigDraftType.DEBUG_BUILD and draft.account_id == user.id: - return AgentConfigDraftType.DEBUG_BUILD.value - if draft.draft_type == AgentConfigDraftType.DRAFT and draft.account_id is None: - return AgentConfigDraftType.DRAFT.value - return AgentConfigDraftType.DRAFT.value + .where( + AgentConfigDraft.tenant_id == app_model.tenant_id, + AgentConfigDraft.draft_type == AgentConfigDraftType.DEBUG_BUILD, + AgentConfigDraft.account_id == user.id, + AgentWorkspaceBinding.tenant_id == app_model.tenant_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + AgentWorkspaceBinding.pending_form_id == form_id, + ) + ) + if build_draft is not None: + return AgentConfigDraftType.DEBUG_BUILD.value, build_draft.id + return AgentConfigDraftType.DRAFT.value, None def _generate_worker( self, @@ -482,7 +499,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): invoke_from=application_generate_entity.invoke_from, ) with session_factory.create_session() as session: - _, _, agent_soul = self._resolve_agent_by_id( + agent, config_version, agent_soul = self._resolve_agent_by_id( tenant_id=app_config.tenant_id, agent_id=application_generate_entity.agent_id, snapshot_id=application_generate_entity.agent_config_snapshot_id, @@ -496,13 +513,18 @@ class AgentAppGenerator(MessageBasedAppGenerator): agent_config_snapshot_id=application_generate_entity.agent_config_snapshot_id, agent_config_version_kind=application_generate_entity.agent_config_version_kind, agent_soul=agent_soul, + home_snapshot_id=config_version.home_snapshot_id, conversation_id=conversation.id, query=query, message_id=message.id, model_name=application_generate_entity.model_conf.model, queue_manager=queue_manager, - session_scope_snapshot_id=application_generate_entity.agent_runtime_session_snapshot_id, - agent_runtime_exit_intent=application_generate_entity.agent_runtime_exit_intent, + session_scope_snapshot_id=application_generate_entity.agent_session_scope_config_version_id, + build_draft_id=( + application_generate_entity.agent_config_snapshot_id + if application_generate_entity.agent_config_version_kind == AgentConfigVersionKind.BUILD_DRAFT + else None + ), ) except GenerateTaskStoppedError: pass @@ -519,18 +541,6 @@ class AgentAppGenerator(MessageBasedAppGenerator): raise AgentAppGeneratorError("query is required") return query.replace("\x00", "") - @staticmethod - def _resolve_agent_runtime_exit_intent(args: Mapping[str, Any]) -> AgentRuntimeExitIntent: - """Resolve API-internal runtime exit policy from controller-owned args. - - Only the private controller-injected "delete" value changes behavior. - Normal chat and resume flows default/fallback to "suspend" so public - payloads and invalid internal values preserve existing semantics. - """ - if args.get(AGENT_RUNTIME_EXIT_INTENT_ARG) == "delete": - return "delete" - return "suspend" - @staticmethod def _build_runner(dify_context: DifyRunContext) -> AgentAppRunner: credentials_provider, _ = build_dify_model_access(dify_context) @@ -546,7 +556,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): stream_run_timeout_seconds=dify_config.AGENT_BACKEND_RUN_TIMEOUT_SECONDS, ), event_adapter=AgentBackendRunEventAdapter(), - session_store=AgentAppRuntimeSessionStore(), + session_store=AgentAppWorkspaceStore(), text_delta_debounce_seconds=dify_config.AGENT_APP_TEXT_DELTA_DEBOUNCE_SECONDS, ) @@ -610,8 +620,10 @@ class AgentAppGenerator(MessageBasedAppGenerator): *, invoke_from: InvokeFrom, draft_type: Any, + draft_id: str | None = None, user: Account | EndUser, session: Session, + conversation: Conversation | None = None, ) -> tuple[Agent, str, Literal["snapshot", "draft", "build_draft"], AgentSoulConfig]: agent = session.scalar( select(Agent) @@ -643,6 +655,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): tenant_id=app_model.tenant_id, agent=agent, draft_type=draft_type, + draft_id=draft_id, account_id=user.id if isinstance(user, Account) else None, session=session, ) @@ -655,28 +668,82 @@ class AgentAppGenerator(MessageBasedAppGenerator): # Public runtime must keep serving the active snapshot even when unpublished draft edits exist. if not agent.active_config_snapshot_id: raise AgentAppNotPublishedError("Agent has not been published") + conversation_binding = self._resolve_conversation_binding( + session=session, + tenant_id=app_model.tenant_id, + app_id=app_model.id, + agent_id=agent.id, + conversation=conversation, + ) + snapshot_id = ( + conversation_binding.agent_config_version_id + if conversation_binding is not None + else agent.active_config_snapshot_id + ) _, snapshot, agent_soul = self._resolve_agent_by_id( tenant_id=app_model.tenant_id, agent_id=agent.id, - snapshot_id=agent.active_config_snapshot_id, + snapshot_id=snapshot_id, session=session, ) + if conversation_binding is not None: + AgentWorkspaceService.validate_binding_generation( + conversation_binding, + base_home_snapshot_id=snapshot.home_snapshot_id, + agent_config_version_id=snapshot.id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) return agent, snapshot.id, "snapshot", agent_soul @staticmethod - def _runtime_session_snapshot_id(*, invoke_from: InvokeFrom, snapshot_id: str) -> str | None: - """Return the session scope snapshot id for Agent App runtime state. + def _resolve_conversation_binding( + *, + session: Session, + tenant_id: str, + app_id: str, + agent_id: str, + conversation: Conversation | None, + ) -> AgentWorkspaceBinding | None: + """Resolve the exact participant generation owned by an existing conversation.""" + + if conversation is None or conversation.agent_workspace_binding_id is None: + return None + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=tenant_id, + binding_id=conversation.agent_workspace_binding_id, + expected_owner_scope=WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=conversation.id, + ), + ) + if binding is None or binding.agent_id != agent_id: + raise AgentAppGeneratorError("Conversation participant Binding is unavailable") + return binding + + @staticmethod + def _session_scope_config_version_id(*, invoke_from: InvokeFrom, config_version_id: str) -> str | None: + """Return the config version id that scopes Agent App session reuse. Console preview/debug chat uses a stable Agent draft row id; build mode uses the current user's build-draft row id. Published/web/API runs use - immutable published snapshot ids. This keeps runtime session continuity + immutable published snapshot ids. This keeps Workspace Binding continuity inside one editable surface without mixing draft/build/published state. """ - return snapshot_id + del invoke_from + return config_version_id @staticmethod def _resolve_debug_draft( - *, tenant_id: str, agent: Agent, draft_type: Any, account_id: str | None, session: Session + *, + tenant_id: str, + agent: Agent, + draft_type: Any, + account_id: str | None, + session: Session, + draft_id: str | None = None, ) -> AgentConfigDraft: effective_draft_type = ( AgentConfigDraftType.DEBUG_BUILD @@ -700,6 +767,8 @@ class AgentAppGenerator(MessageBasedAppGenerator): AgentConfigDraft.draft_type == AgentConfigDraftType.DEBUG_BUILD, AgentConfigDraft.account_id == account_id, ) + if draft_id is not None: + stmt = stmt.where(AgentConfigDraft.id == draft_id) draft = session.scalar(stmt.order_by(AgentConfigDraft.updated_at.desc()).limit(1)) if draft is not None: return draft diff --git a/api/core/app/apps/agent_app/app_runner.py b/api/core/app/apps/agent_app/app_runner.py index f844068073d..6202677c332 100644 --- a/api/core/app/apps/agent_app/app_runner.py +++ b/api/core/app/apps/agent_app/app_runner.py @@ -3,8 +3,7 @@ Unlike the legacy ``AgentChatAppRunner`` (which runs an in-process ReAct loop), this runner delegates to the Agent backend, consumes the streamed event flow, republishes the assistant answer through the existing EasyUI chat task -pipeline, and then either saves or retires the conversation-owned runtime -session depending on the turn's exit policy. +pipeline, and saves the latest Agenton snapshot on the persistent Binding. """ from __future__ import annotations @@ -31,22 +30,20 @@ from clients.agent_backend import ( AgentBackendRunFailedInternalEvent, AgentBackendRunSucceededInternalEvent, AgentBackendStreamInternalEvent, - extract_runtime_layer_specs, ) -from clients.agent_backend.session_cleanup import AgentBackendSessionCleanupPayload from core.app.apps.agent_app.runtime_request_builder import ( AgentAppRuntimeBuildContext, AgentAppRuntimeRequest, AgentAppRuntimeRequestBuilder, ) from core.app.apps.agent_app.session_store import ( - AgentAppRuntimeSessionStore, AgentAppSessionScope, + AgentAppWorkspaceStore, StoredAgentAppSession, ) from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import AgentRuntimeExitIntent, DifyRunContext +from core.app.entities.app_invoke_entities import DifyRunContext from core.app.entities.queue_entities import ( QueueAgentMessageEvent, QueueAgentThoughtEvent, @@ -67,10 +64,10 @@ from graphon.model_runtime.errors.invoke import ( InvokeRateLimitError, InvokeServerUnavailableError, ) +from models.agent import AgentConfigVersionKind from models.agent_config_entities import AgentSoulConfig from models.enums import CreatorUserRole from models.model import MessageAgentThought -from tasks.agent_backend_session_cleanup_task import cleanup_conversation_agent_runtime_session logger = logging.getLogger(__name__) @@ -620,7 +617,7 @@ class AgentAppRunner: request_builder: AgentAppRuntimeRequestBuilder, agent_backend_client: AgentBackendRunClient, event_adapter: AgentBackendRunEventAdapter, - session_store: AgentAppRuntimeSessionStore, + session_store: AgentAppWorkspaceStore, text_delta_debounce_seconds: float, ) -> None: self._request_builder = request_builder @@ -637,37 +634,41 @@ class AgentAppRunner: agent_config_snapshot_id: str, agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot", agent_soul: AgentSoulConfig, + home_snapshot_id: str, conversation_id: str, query: str, message_id: str, model_name: str, queue_manager: AppQueueManager, session_scope_snapshot_id: str | None | _DefaultSessionScopeSnapshotId = _DEFAULT_SESSION_SCOPE_SNAPSHOT_ID, - agent_runtime_exit_intent: AgentRuntimeExitIntent = "suspend", + build_draft_id: str | None = None, ) -> None: - preserve_session = agent_runtime_exit_intent == "suspend" scope = self._build_session_scope( dify_context=dify_context, agent_id=agent_id, agent_config_snapshot_id=agent_config_snapshot_id, + home_snapshot_id=home_snapshot_id, conversation_id=conversation_id, session_scope_snapshot_id=session_scope_snapshot_id, + agent_config_version_kind=AgentConfigVersionKind(agent_config_version_kind), + build_draft_id=build_draft_id, ) # ENG-638: if a prior turn paused on ask_human and the form is now answered, # resume by threading the human's reply into this run as deferred_tool_results. - stored = self._session_store.load_active_session(scope) + stored = self._session_store.load_or_create(scope) runtime = self._build_runtime( dify_context=dify_context, agent_id=agent_id, agent_config_snapshot_id=agent_config_snapshot_id, agent_config_version_kind=agent_config_version_kind, agent_soul=agent_soul, + binding_id=stored.binding_id, + backend_binding_ref=stored.backend_binding_ref, conversation_id=conversation_id, query=query, idempotency_key=message_id, stored=stored, message_id=message_id, - suspend_on_exit=preserve_session, ) create_response = self._agent_backend_client.create_run(runtime.request) @@ -681,9 +682,6 @@ class AgentAppRunner: ) if isinstance(terminal, AgentBackendDeferredToolCallInternalEvent): - if not preserve_session: - self._mark_session_cleaned(scope=scope, backend_run_id=terminal.run_id) - raise AgentBackendError("Agent App finalization cannot pause for human input.") # ENG-635: the agent asked a human. End this turn with the question and # a conversation-owned HITL form; a form submission resumes the run. self._pause_for_ask_human( @@ -703,8 +701,8 @@ class AgentAppRunner: if not isinstance(terminal, AgentBackendRunSucceededInternalEvent): if isinstance(terminal, AgentBackendRunFailedInternalEvent): reason = terminal.reason - if reason == "sandbox_expired": - raise AgentBackendError("The agent session sandbox has expired. Please start a new conversation.") + if reason == "binding_lost": + raise AgentBackendError("The retained agent working environment is no longer available.") raise _agent_backend_failure_to_exception(terminal) raise AgentBackendError("Agent backend run did not complete successfully.") @@ -719,38 +717,18 @@ class AgentAppRunner: message_id, exc_info=True, ) - if preserve_session: - superseded_sessions = self._load_superseded_sessions(scope=scope) - self._publish_terminal_answer( - queue_manager=queue_manager, - model_name=model_name, - answer=answer, - query=query, - usage=_llm_usage_from_agent_backend(terminal.usage), - ) - session_saved = self._save_session( - scope=scope, - backend_run_id=terminal.run_id, - snapshot=terminal.session_snapshot, - runtime_layer_specs=extract_runtime_layer_specs(runtime.request.composition), - ) - if session_saved: - self._cleanup_superseded_sessions(superseded_sessions) - else: - # The backend has already accepted a terminal success with - # delete-on-exit semantics. Local publish/persistence errors must - # not keep the API-side session row active, and cleanup failures - # must not replace the original publish/error outcome. - try: - self._publish_terminal_answer( - queue_manager=queue_manager, - model_name=model_name, - answer=answer, - query=query, - usage=_llm_usage_from_agent_backend(terminal.usage), - ) - finally: - self._mark_session_cleaned(scope=scope, backend_run_id=terminal.run_id) + self._publish_terminal_answer( + queue_manager=queue_manager, + model_name=model_name, + answer=answer, + query=query, + usage=_llm_usage_from_agent_backend(terminal.usage), + ) + self._save_session( + scope=scope, + binding_id=runtime.binding_id, + snapshot=terminal.session_snapshot, + ) def _build_session_scope( self, @@ -758,8 +736,11 @@ class AgentAppRunner: dify_context: DifyRunContext, agent_id: str, agent_config_snapshot_id: str, + home_snapshot_id: str, conversation_id: str, session_scope_snapshot_id: str | None | _DefaultSessionScopeSnapshotId, + agent_config_version_kind: AgentConfigVersionKind, + build_draft_id: str | None = None, ) -> AgentAppSessionScope: if isinstance(session_scope_snapshot_id, _DefaultSessionScopeSnapshotId): effective_session_scope_snapshot_id: str | None = agent_config_snapshot_id @@ -770,7 +751,10 @@ class AgentAppRunner: app_id=dify_context.app_id, conversation_id=conversation_id, agent_id=agent_id, - agent_config_snapshot_id=effective_session_scope_snapshot_id, + agent_config_snapshot_id=effective_session_scope_snapshot_id or agent_config_snapshot_id, + home_snapshot_id=home_snapshot_id, + agent_config_version_kind=agent_config_version_kind, + build_draft_id=build_draft_id, ) def _build_runtime( @@ -781,14 +765,15 @@ class AgentAppRunner: agent_config_snapshot_id: str, agent_config_version_kind: Literal["snapshot", "draft", "build_draft"], agent_soul: AgentSoulConfig, + binding_id: str, + backend_binding_ref: str, conversation_id: str, query: str, idempotency_key: str, - stored: StoredAgentAppSession | None, + stored: StoredAgentAppSession, message_id: str | None, - suspend_on_exit: bool, ) -> AgentAppRuntimeRequest: - session_snapshot = stored.session_snapshot if stored is not None else None + session_snapshot = stored.session_snapshot deferred_tool_results = ( self._resolve_pending_ask_human(stored=stored, dify_context=dify_context, message_id=message_id) if message_id is not None @@ -804,9 +789,10 @@ class AgentAppRunner: conversation_id=conversation_id, user_query=query, idempotency_key=idempotency_key, + binding_id=binding_id, + backend_binding_ref=backend_binding_ref, session_snapshot=session_snapshot, deferred_tool_results=deferred_tool_results, - suspend_on_exit=suspend_on_exit, ) ) @@ -843,9 +829,8 @@ class AgentAppRunner: # second run with the human's answer (ENG-637/638 columns, conversation owner). self._save_session( scope=scope, - backend_run_id=terminal.run_id, + binding_id=runtime.binding_id, snapshot=terminal.session_snapshot, - runtime_layer_specs=extract_runtime_layer_specs(runtime.request.composition), pending_form_id=created.form_id, pending_tool_call_id=terminal.deferred_tool_call.tool_call_id, ) @@ -862,12 +847,12 @@ class AgentAppRunner: def _resolve_pending_ask_human( self, *, - stored: StoredAgentAppSession | None, + stored: StoredAgentAppSession, dify_context: DifyRunContext, message_id: str, ) -> DeferredToolResultsPayload | None: """Build deferred_tool_results when a pending ask_human form is answered.""" - if stored is None or stored.pending_form_id is None or stored.pending_tool_call_id is None: + if stored.pending_form_id is None or stored.pending_tool_call_id is None: return None outcome = resolve_ask_human_form( form_id=stored.pending_form_id, @@ -1036,18 +1021,16 @@ class AgentAppRunner: self, *, scope: AgentAppSessionScope, - backend_run_id: str, + binding_id: str, snapshot: Any, - runtime_layer_specs: Any, pending_form_id: str | None = None, pending_tool_call_id: str | None = None, ) -> bool: try: self._session_store.save_active_snapshot( scope=scope, - backend_run_id=backend_run_id, + binding_id=binding_id, snapshot=snapshot, - runtime_layer_specs=runtime_layer_specs, pending_form_id=pending_form_id, pending_tool_call_id=pending_tool_call_id, ) @@ -1064,87 +1047,6 @@ class AgentAppRunner: ) return False - def _load_superseded_sessions(self, *, scope: AgentAppSessionScope) -> list[StoredAgentAppSession]: - try: - stored_sessions = self._session_store.list_active_sessions_for_conversation( - tenant_id=scope.tenant_id, - app_id=scope.app_id, - conversation_id=scope.conversation_id, - ) - except Exception: - logger.warning( - "Failed to load existing Agent App conversation sessions before snapshot save: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s", - scope.tenant_id, - scope.app_id, - scope.conversation_id, - scope.agent_id, - exc_info=True, - ) - return [] - - return [stored for stored in stored_sessions if stored.scope != scope] - - def _cleanup_superseded_sessions(self, stored_sessions: list[StoredAgentAppSession]) -> None: - for stored_session in stored_sessions: - try: - if stored_session.runtime_layer_specs: - payload = AgentBackendSessionCleanupPayload( - session_snapshot=stored_session.session_snapshot, - runtime_layer_specs=stored_session.runtime_layer_specs, - idempotency_key=( - f"{stored_session.scope.tenant_id}:{stored_session.scope.app_id}:" - f"{stored_session.scope.conversation_id}:{stored_session.scope.agent_id}:" - f"{stored_session.scope.agent_config_snapshot_id or 'no-config'}:" - f"superseded-session-cleanup:{stored_session.backend_run_id or 'no-run'}" - ), - metadata={ - "tenant_id": stored_session.scope.tenant_id, - "app_id": stored_session.scope.app_id, - "conversation_id": stored_session.scope.conversation_id, - "agent_id": stored_session.scope.agent_id, - "agent_config_snapshot_id": stored_session.scope.agent_config_snapshot_id, - "previous_agent_backend_run_id": stored_session.backend_run_id, - }, - ) - cleanup_conversation_agent_runtime_session.delay(payload.model_dump(mode="json")) - except Exception: - logger.warning( - "Failed to enqueue Agent backend cleanup for superseded Agent App session: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", - stored_session.scope.tenant_id, - stored_session.scope.app_id, - stored_session.scope.conversation_id, - stored_session.scope.agent_id, - stored_session.backend_run_id, - exc_info=True, - ) - - def _mark_session_cleaned( - self, - *, - scope: AgentAppSessionScope, - backend_run_id: str, - ) -> None: - """Best-effort delete-on-exit cleanup for the API-side session row. - - Once the Agent backend reaches a terminal event, cleanup persistence - must not replace the original publish/error outcome for that turn. - """ - try: - self._session_store.mark_cleaned(scope=scope, backend_run_id=backend_run_id) - except Exception: - logger.warning( - "Failed to retire Agent App conversation session after delete-on-exit: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", - scope.tenant_id, - scope.app_id, - scope.conversation_id, - scope.agent_id, - backend_run_id, - exc_info=True, - ) - @staticmethod def _terminal_output_to_answer(output: JsonValue) -> str: """Normalize the backend's terminal output to assistant text. diff --git a/api/core/app/apps/agent_app/runtime_request_builder.py b/api/core/app/apps/agent_app/runtime_request_builder.py index c90ab77d681..ccad454d8a8 100644 --- a/api/core/app/apps/agent_app/runtime_request_builder.py +++ b/api/core/app/apps/agent_app/runtime_request_builder.py @@ -70,11 +70,12 @@ class AgentAppRuntimeBuildContext: conversation_id: str user_query: str idempotency_key: str + binding_id: str + backend_binding_ref: str agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot" session_snapshot: CompositorSessionSnapshot | None = None # ENG-638: set when resuming a chat turn after a submitted ask_human form. deferred_tool_results: DeferredToolResultsPayload | None = None - suspend_on_exit: bool = True @dataclass(frozen=True, slots=True) @@ -82,6 +83,7 @@ class AgentAppRuntimeRequest: request: CreateRunRequest redacted_request: dict[str, Any] metadata: dict[str, Any] + binding_id: str class AgentAppRuntimeRequestBuilder: @@ -160,6 +162,7 @@ class AgentAppRuntimeRequestBuilder: invoke_from=cast(DifyExecutionContextInvokeFrom, context.dify_context.invoke_from.value), agent_mode="agent_app", ), + backend_binding_ref=context.backend_binding_ref, # ENG-616: expand slash-menu mention tokens to canonical names so # no frontend-internal {{#…#}} marker ever reaches the model. agent_soul_prompt=expand_prompt_mentions(agent_soul.prompt.system_prompt, soul_prompt_resolver).strip() @@ -175,13 +178,17 @@ class AgentAppRuntimeRequestBuilder: shell_config=build_shell_layer_config(agent_soul), session_snapshot=context.session_snapshot, deferred_tool_results=context.deferred_tool_results, - suspend_on_exit=context.suspend_on_exit, idempotency_key=context.idempotency_key, metadata=metadata, ) ) redacted = cast(dict[str, Any], redact_for_agent_backend_log(request)) - return AgentAppRuntimeRequest(request=request, redacted_request=redacted, metadata=metadata) + return AgentAppRuntimeRequest( + request=request, + redacted_request=redacted, + metadata=metadata, + binding_id=context.binding_id, + ) def _build_tool_layers( self, diff --git a/api/core/app/apps/agent_app/session_store.py b/api/core/app/apps/agent_app/session_store.py index 7696155a1c4..c07c309aec6 100644 --- a/api/core/app/apps/agent_app/session_store.py +++ b/api/core/app/apps/agent_app/session_store.py @@ -1,255 +1,176 @@ -"""Conversation-keyed Agent backend session store for the Agent App type. - -Shares the unified ``agent_runtime_sessions`` table with the workflow Agent -Node store, but owns rows with ``owner_type = conversation``: one Agent App -conversation maps to one Agent session, so multi-turn chat re-enters the same -``session_snapshot``. Cross-conversation memory (PRD Global / Per app) is a -phase-2 concern and not modeled here. -""" +"""Persist and resolve the exact participant owned by an Agent App caller.""" from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from agenton.compositor import CompositorSessionSnapshot -from dify_agent.protocol import RuntimeLayerSpec -from pydantic import TypeAdapter from sqlalchemy import select +from sqlalchemy.orm import Session from core.db.session_factory import session_factory -from libs.datetime_utils import naive_utc_now from models.agent import ( - AgentRuntimeSession, - AgentRuntimeSessionOwnerType, - AgentRuntimeSessionStatus, + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigVersionKind, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) +from models.model import App, Conversation +from services.agent.workspace_service import ( + AgentWorkspaceNotFoundError, + AgentWorkspaceService, + WorkspaceOwnerScope, ) - -_RUNTIME_LAYER_SPECS_ADAPTER: TypeAdapter[list[RuntimeLayerSpec]] = TypeAdapter(list[RuntimeLayerSpec]) - - -def _serialize_runtime_layer_specs(specs: list[RuntimeLayerSpec]) -> str: - return _RUNTIME_LAYER_SPECS_ADAPTER.dump_json(specs).decode() - - -def _deserialize_runtime_layer_specs(value: str | None) -> list[RuntimeLayerSpec]: - if not value: - return [] - return _RUNTIME_LAYER_SPECS_ADAPTER.validate_json(value) @dataclass(frozen=True, slots=True) class AgentAppSessionScope: - """Identity of one Agent App conversation session.""" - tenant_id: str app_id: str conversation_id: str agent_id: str - agent_config_snapshot_id: str | None + agent_config_snapshot_id: str + home_snapshot_id: str + agent_config_version_kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT + build_draft_id: str | None = None + + @property + def workspace_owner(self) -> WorkspaceOwnerScope: + owner_type = ( + AgentWorkspaceOwnerType.BUILD_DRAFT if self.build_draft_id else AgentWorkspaceOwnerType.CONVERSATION + ) + return WorkspaceOwnerScope( + tenant_id=self.tenant_id, + app_id=self.app_id, + owner_type=owner_type, + owner_id=self.build_draft_id or self.conversation_id, + ) @dataclass(frozen=True, slots=True) class StoredAgentAppSession: - """Persisted Agent App conversation session with reusable runtime specs.""" - scope: AgentAppSessionScope - session_snapshot: CompositorSessionSnapshot - backend_run_id: str | None - runtime_layer_specs: list[RuntimeLayerSpec] = field(default_factory=list) - # ENG-635: set while the conversation turn is paused on a dify.ask_human - # deferred call, awaiting a HITL form submission. + binding_id: str + workspace_id: str + backend_binding_ref: str + session_snapshot: CompositorSessionSnapshot | None pending_form_id: str | None = None pending_tool_call_id: str | None = None -class AgentAppRuntimeSessionStore: - """Persists Agent backend session snapshots for Agent App conversations.""" +class AgentAppWorkspaceStore: + """Resolve Agent App sessions through a caller-owned Binding pointer.""" - def load_active_snapshot(self, scope: AgentAppSessionScope) -> CompositorSessionSnapshot | None: - stored = self.load_active_session(scope) - return stored.session_snapshot if stored is not None else None - - def load_active_session(self, scope: AgentAppSessionScope) -> StoredAgentAppSession | None: + def load_or_create(self, scope: AgentAppSessionScope) -> StoredAgentAppSession: with session_factory.create_session() as session: - row = session.scalar(self._active_stmt(scope)) - if row is None: - return None - return StoredAgentAppSession( - scope=scope, - session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), - backend_run_id=row.backend_run_id, - runtime_layer_specs=_deserialize_runtime_layer_specs(row.composition_layer_specs), - pending_form_id=row.pending_form_id, - pending_tool_call_id=row.pending_tool_call_id, - ) - - def load_active_session_for_conversation( - self, *, tenant_id: str, app_id: str, conversation_id: str - ) -> StoredAgentAppSession | None: - """Load the latest ACTIVE session for one conversation-level sandbox lookup. - - Sandbox inspection only knows the product locator - ``tenant_id + app_id + conversation_id``; it does not know which - ``agent_id`` or Agent Soul snapshot produced the active shell session. - This method therefore resolves the newest ACTIVE conversation-owned row - for that conversation and returns both the resumable snapshot and the - persisted non-sensitive runtime layer specs needed to build a - ``SandboxLocator``. - """ - stmt = ( - select(AgentRuntimeSession) - .where( - AgentRuntimeSession.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION, - AgentRuntimeSession.tenant_id == tenant_id, - AgentRuntimeSession.app_id == app_id, - AgentRuntimeSession.conversation_id == conversation_id, - AgentRuntimeSession.status == AgentRuntimeSessionStatus.ACTIVE, - ) - .order_by(AgentRuntimeSession.updated_at.desc()) - ) - with session_factory.create_session() as session: - row = session.scalar(stmt) - if row is None: - return None - return StoredAgentAppSession( - scope=AgentAppSessionScope( - tenant_id=row.tenant_id, - app_id=row.app_id, - conversation_id=row.conversation_id or "", - agent_id=row.agent_id, - agent_config_snapshot_id=row.agent_config_snapshot_id or "", - ), - session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), - backend_run_id=row.backend_run_id, - runtime_layer_specs=_deserialize_runtime_layer_specs(row.composition_layer_specs), - ) - - def list_active_sessions_for_conversation( - self, *, tenant_id: str, app_id: str, conversation_id: str - ) -> list[StoredAgentAppSession]: - """List all ACTIVE conversation-owned sessions for lifecycle cleanup.""" - stmt = ( - select(AgentRuntimeSession) - .where( - AgentRuntimeSession.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION, - AgentRuntimeSession.tenant_id == tenant_id, - AgentRuntimeSession.app_id == app_id, - AgentRuntimeSession.conversation_id == conversation_id, - AgentRuntimeSession.status == AgentRuntimeSessionStatus.ACTIVE, - ) - .order_by(AgentRuntimeSession.updated_at.desc()) - ) - with session_factory.create_session() as session: - rows = session.scalars(stmt).all() - return [ - StoredAgentAppSession( - scope=AgentAppSessionScope( - tenant_id=row.tenant_id, - app_id=row.app_id, - conversation_id=row.conversation_id or "", - agent_id=row.agent_id, - agent_config_snapshot_id=row.agent_config_snapshot_id, - ), - session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), - backend_run_id=row.backend_run_id, - runtime_layer_specs=_deserialize_runtime_layer_specs(row.composition_layer_specs), - pending_form_id=row.pending_form_id, - pending_tool_call_id=row.pending_tool_call_id, + caller = self._load_caller(session=session, scope=scope) + binding_id = caller.agent_workspace_binding_id + if binding_id is None: + binding = AgentWorkspaceService.create_binding( + session=session, + scope=scope.workspace_owner, + agent_id=scope.agent_id, + base_home_snapshot_id=scope.home_snapshot_id, + agent_config_version_id=scope.agent_config_snapshot_id, + agent_config_version_kind=scope.agent_config_version_kind, ) - for row in rows - ] + caller.agent_workspace_binding_id = binding.id + session.commit() + else: + binding = self._get_binding(session=session, scope=scope, binding_id=binding_id) + return self._stored(scope, binding) + + @staticmethod + def _load_caller(*, session: Session, scope: AgentAppSessionScope) -> Conversation | AgentConfigDraft: + if scope.build_draft_id is not None: + if scope.agent_config_version_kind != AgentConfigVersionKind.BUILD_DRAFT: + raise AgentWorkspaceNotFoundError("Build Draft caller requires build_draft generation") + draft = session.scalar( + select(AgentConfigDraft).where( + AgentConfigDraft.id == scope.build_draft_id, + AgentConfigDraft.tenant_id == scope.tenant_id, + AgentConfigDraft.agent_id == scope.agent_id, + AgentConfigDraft.draft_type == AgentConfigDraftType.DEBUG_BUILD, + ) + ) + if draft is None: + raise AgentWorkspaceNotFoundError("Build Draft caller is unavailable") + return draft + if scope.agent_config_version_kind == AgentConfigVersionKind.BUILD_DRAFT: + raise AgentWorkspaceNotFoundError("Build Draft caller ID is required") + conversation = session.scalar( + select(Conversation) + .join(App, App.id == Conversation.app_id) + .where( + App.tenant_id == scope.tenant_id, + Conversation.id == scope.conversation_id, + Conversation.app_id == scope.app_id, + Conversation.is_deleted.is_(False), + ) + ) + if conversation is None: + raise AgentWorkspaceNotFoundError("Conversation caller is unavailable") + return conversation + + @staticmethod + def _get_binding( + *, + session: Session, + scope: AgentAppSessionScope, + binding_id: str, + ) -> AgentWorkspaceBinding: + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=scope.tenant_id, + binding_id=binding_id, + expected_owner_scope=scope.workspace_owner, + ) + if binding is None or binding.agent_id != scope.agent_id: + raise AgentWorkspaceNotFoundError("Caller participant Binding is unavailable") + AgentWorkspaceService.validate_binding_generation( + binding, + base_home_snapshot_id=scope.home_snapshot_id, + agent_config_version_id=scope.agent_config_snapshot_id, + agent_config_version_kind=scope.agent_config_version_kind, + ) + return binding def save_active_snapshot( self, *, scope: AgentAppSessionScope, - backend_run_id: str, + binding_id: str, snapshot: CompositorSessionSnapshot | None, - runtime_layer_specs: list[RuntimeLayerSpec], pending_form_id: str | None = None, pending_tool_call_id: str | None = None, ) -> None: - """Persist the current conversation snapshot and enforce one ACTIVE row. - - Agent App chat treats one conversation as one resumable runtime shell. - Saving the latest snapshot therefore upserts the scoped row back to - ACTIVE and retires any other ACTIVE conversation-owned rows for the - same ``tenant_id + app_id + conversation_id`` so later lookups see a - single active session. - """ if snapshot is None: return - snapshot_json = snapshot.model_dump_json() - runtime_layer_specs_json = _serialize_runtime_layer_specs(runtime_layer_specs) - with session_factory.create_session() as session: - row = session.scalar(self._scope_stmt(scope)) - if row is None: - row = AgentRuntimeSession( - tenant_id=scope.tenant_id, - app_id=scope.app_id, - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id=scope.agent_id, - agent_config_snapshot_id=scope.agent_config_snapshot_id, - conversation_id=scope.conversation_id, - backend_run_id=backend_run_id, - session_snapshot=snapshot_json, - composition_layer_specs=runtime_layer_specs_json, - status=AgentRuntimeSessionStatus.ACTIVE, - pending_form_id=pending_form_id, - pending_tool_call_id=pending_tool_call_id, - ) - session.add(row) - else: - row.backend_run_id = backend_run_id - row.session_snapshot = snapshot_json - row.composition_layer_specs = runtime_layer_specs_json - row.status = AgentRuntimeSessionStatus.ACTIVE - row.cleaned_at = None - # Set (or clear, when omitted) the ask_human pause correlation. - row.pending_form_id = pending_form_id - row.pending_tool_call_id = pending_tool_call_id - session.flush() - other_rows = session.scalars( - select(AgentRuntimeSession).where( - AgentRuntimeSession.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION, - AgentRuntimeSession.tenant_id == scope.tenant_id, - AgentRuntimeSession.app_id == scope.app_id, - AgentRuntimeSession.conversation_id == scope.conversation_id, - AgentRuntimeSession.status == AgentRuntimeSessionStatus.ACTIVE, - AgentRuntimeSession.id != row.id, - ) - ).all() - for other_row in other_rows: - other_row.status = AgentRuntimeSessionStatus.CLEANED - other_row.cleaned_at = naive_utc_now() - session.commit() - - def mark_cleaned(self, *, scope: AgentAppSessionScope, backend_run_id: str | None = None) -> None: - with session_factory.create_session() as session: - row = session.scalar(self._active_stmt(scope)) - if row is None: - return - if backend_run_id is not None: - row.backend_run_id = backend_run_id - row.status = AgentRuntimeSessionStatus.CLEANED - row.cleaned_at = naive_utc_now() - session.commit() + AgentWorkspaceService.save_binding_session_snapshot( + tenant_id=scope.tenant_id, + binding_id=binding_id, + session_snapshot=snapshot.model_dump_json(), + pending_form_id=pending_form_id, + pending_tool_call_id=pending_tool_call_id, + ) @staticmethod - def _scope_stmt(scope: AgentAppSessionScope): - stmt = select(AgentRuntimeSession).where( - AgentRuntimeSession.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION, - AgentRuntimeSession.tenant_id == scope.tenant_id, - AgentRuntimeSession.conversation_id == scope.conversation_id, - AgentRuntimeSession.agent_id == scope.agent_id, + def _stored(scope: AgentAppSessionScope, binding: AgentWorkspaceBinding) -> StoredAgentAppSession: + snapshot = ( + CompositorSessionSnapshot.model_validate_json(binding.session_snapshot) + if binding.session_snapshot + else None + ) + return StoredAgentAppSession( + scope=scope, + binding_id=binding.id, + workspace_id=binding.workspace_id, + backend_binding_ref=binding.backend_binding_ref, + session_snapshot=snapshot, + pending_form_id=binding.pending_form_id, + pending_tool_call_id=binding.pending_tool_call_id, ) - if scope.agent_config_snapshot_id is None: - return stmt.where(AgentRuntimeSession.agent_config_snapshot_id.is_(None)) - return stmt.where(AgentRuntimeSession.agent_config_snapshot_id == scope.agent_config_snapshot_id) - - @classmethod - def _active_stmt(cls, scope: AgentAppSessionScope): - return cls._scope_stmt(scope).where(AgentRuntimeSession.status == AgentRuntimeSessionStatus.ACTIVE) -__all__ = ["AgentAppRuntimeSessionStore", "AgentAppSessionScope", "StoredAgentAppSession"] +__all__ = ["AgentAppSessionScope", "AgentAppWorkspaceStore", "StoredAgentAppSession"] diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py index 95c9d777ebd..d7427408792 100644 --- a/api/core/app/apps/workflow/app_runner.py +++ b/api/core/app/apps/workflow/app_runner.py @@ -10,11 +10,11 @@ from core.app.apps.workflow.command_channels import ( CombinedCommandChannel, ) from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity +from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, WorkflowAppGenerateEntity from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.workflow.node_factory import get_default_root_node_id -from core.workflow.nodes.agent_v2.session_cleanup_layer import build_workflow_agent_session_cleanup_layer +from core.workflow.nodes.agent_v2.workspace_retirement_layer import build_workflow_agent_workspace_retirement_layer from core.workflow.snippet_start import get_compatible_start_aliases from core.workflow.system_variables import build_bootstrap_variables, build_system_variables from core.workflow.variable_pool_initializer import add_node_inputs_to_pool, add_variables_to_pool @@ -197,7 +197,18 @@ class WorkflowAppRunner(WorkflowBasedAppRunner): ) workflow_entry.graph_engine.layer(persistence_layer) - workflow_entry.graph_engine.layer(build_workflow_agent_session_cleanup_layer()) + workflow_entry.graph_engine.layer( + build_workflow_agent_workspace_retirement_layer( + dify_run_context=DifyRunContext( + tenant_id=self._workflow.tenant_id, + app_id=self._workflow.app_id, + user_id=self.application_generate_entity.user_id, + user_from=user_from, + invoke_from=invoke_from, + trace_session_id=self.application_generate_entity.extras.get("trace_session_id"), + ) + ) + ) for layer in self._graph_engine_layers: workflow_entry.graph_engine.layer(layer) diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index b5f7515cd9f..b4e2ced08aa 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -15,8 +15,6 @@ if TYPE_CHECKING: DIFY_RUN_CONTEXT_KEY = "_dify" -AGENT_RUNTIME_EXIT_INTENT_ARG = "_agent_runtime_exit_intent" -type AgentRuntimeExitIntent = Literal["suspend", "delete"] class UserFrom(StrEnum): @@ -227,12 +225,8 @@ class AgentAppGenerateEntity(ChatAppGenerateEntity): backend should read from: immutable snapshot, shared draft, or per-user build draft. - ``agent_runtime_session_snapshot_id`` carries the runtime session scope - used to resume or suspend within the same editable config surface. - - ``agent_runtime_exit_intent`` is API-internal lifecycle policy for the - Agent backend session after this turn finishes. Normal chat/resume turns - suspend on exit; build-chat finalization deletes the backend runtime. + ``agent_session_scope_config_version_id`` identifies the draft or immutable + config version whose Workspace Binding should be reused for this session. ``prompt_file_mappings`` preserves the raw request ``files`` array for the Agent backend prompt. These references are appended to the backend prompt @@ -242,8 +236,7 @@ class AgentAppGenerateEntity(ChatAppGenerateEntity): agent_id: str agent_config_snapshot_id: str agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot" - agent_runtime_session_snapshot_id: str | None = None - agent_runtime_exit_intent: AgentRuntimeExitIntent = "suspend" + agent_session_scope_config_version_id: str | None = None prompt_file_mappings: Sequence[JsonValue] = Field(default_factory=list) diff --git a/api/core/app/workflow/layers/persistence.py b/api/core/app/workflow/layers/persistence.py index 52887765c02..415c8a1826d 100644 --- a/api/core/app/workflow/layers/persistence.py +++ b/api/core/app/workflow/layers/persistence.py @@ -20,11 +20,13 @@ from core.helper.trace_id_helper import ParentTraceContext from core.ops.entities.trace_entity import TraceTaskName from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository +from core.workflow.node_execution_process_data import preserve_workflow_agent_binding_id from core.workflow.system_variables import SystemVariableKey from core.workflow.variable_prefixes import SYSTEM_VARIABLE_NODE_ID from core.workflow.workflow_run_outputs import project_node_outputs_for_workflow_run from graphon.entities import WorkflowExecution, WorkflowNodeExecution from graphon.enums import ( + BuiltinNodeTypes, WorkflowExecutionStatus, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus, @@ -241,7 +243,10 @@ class WorkflowPersistenceLayer(GraphEngineLayer): ) self._node_execution_cache[event.id] = domain_execution - self._workflow_node_execution_repository.save(domain_execution) + if event.node_type == BuiltinNodeTypes.AGENT and event.node_version == "2": + self._workflow_node_execution_repository.save_synchronously(domain_execution) + else: + self._workflow_node_execution_repository.save(domain_execution) snapshot = _NodeRuntimeSnapshot( node_id=event.node_id, @@ -361,7 +366,11 @@ class WorkflowPersistenceLayer(GraphEngineLayer): def _append_retry_history(self, execution: WorkflowNodeExecution, event: NodeRunRetryEvent) -> None: """Append a validated full attempt before repository truncation or offload.""" finished_at = naive_utc_now() - process_data = dict(execution.process_data or {}) + process_data = preserve_workflow_agent_binding_id( + event.node_run_result.process_data, + execution.process_data, + ) + process_data = dict(process_data or {}) raw_history = process_data.get(RETRY_HISTORY_PROCESS_DATA_KEY) history = list(raw_history) if isinstance(raw_history, list) else [] projected_outputs = project_node_outputs_for_workflow_run( @@ -390,11 +399,12 @@ class WorkflowPersistenceLayer(GraphEngineLayer): next_process_data: Mapping[str, Any] | None, ) -> Mapping[str, Any] | None: """Keep internal retry history while replacing node-specific Process Data.""" + merged_process_data = preserve_workflow_agent_binding_id(existing_process_data, next_process_data) raw_history = (existing_process_data or {}).get(RETRY_HISTORY_PROCESS_DATA_KEY) if not isinstance(raw_history, list) or not raw_history: - return next_process_data + return merged_process_data - merged_process_data = dict(next_process_data or {}) + merged_process_data = dict(merged_process_data or {}) merged_process_data[RETRY_HISTORY_PROCESS_DATA_KEY] = raw_history return merged_process_data @@ -440,6 +450,11 @@ class WorkflowPersistenceLayer(GraphEngineLayer): outputs=projected_outputs, metadata=node_result.metadata, ) + else: + domain_execution.process_data = preserve_workflow_agent_binding_id( + node_result.process_data, + domain_execution.process_data, + ) self._workflow_node_execution_repository.save(domain_execution) self._workflow_node_execution_repository.save_execution_data(domain_execution) diff --git a/api/core/repositories/celery_workflow_node_execution_repository.py b/api/core/repositories/celery_workflow_node_execution_repository.py index 3b152a0e232..697429eccdc 100644 --- a/api/core/repositories/celery_workflow_node_execution_repository.py +++ b/api/core/repositories/celery_workflow_node_execution_repository.py @@ -16,6 +16,9 @@ from core.repositories.factory import ( OrderConfig, WorkflowNodeExecutionRepository, ) +from core.repositories.sqlalchemy_workflow_node_execution_repository import ( + SQLAlchemyWorkflowNodeExecutionRepository, +) from graphon.entities import WorkflowNodeExecution from models import Account, CreatorUserRole, EndUser from models.workflow import WorkflowNodeExecutionTriggeredFrom @@ -49,6 +52,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): _creator_user_role: CreatorUserRole _execution_cache: dict[str, WorkflowNodeExecution] _workflow_execution_mapping: dict[str, list[str]] + _sql_repository: SQLAlchemyWorkflowNodeExecutionRepository def __init__( self, @@ -98,6 +102,13 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): # Cache for mapping workflow_execution_ids to execution IDs for efficient retrieval self._workflow_execution_mapping = {} + self._sql_repository = SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=session_factory, + tenant_id=tenant_id, + user=user, + app_id=app_id, + triggered_from=triggered_from, + ) logger.info( "Initialized CeleryWorkflowNodeExecutionRepository for tenant %s, app %s, triggered_from %s", @@ -149,6 +160,17 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): # For now, we'll re-raise the exception raise + @override + def save_synchronously(self, execution: WorkflowNodeExecution) -> None: + """Create the Agent v2 caller row before runtime participant allocation.""" + + self._sql_repository.save_synchronously(execution) + self._execution_cache[execution.id] = execution + if execution.workflow_execution_id: + execution_ids = self._workflow_execution_mapping.setdefault(execution.workflow_execution_id, []) + if execution.id not in execution_ids: + execution_ids.append(execution.id) + @override def get_by_workflow_execution( self, diff --git a/api/core/repositories/factory.py b/api/core/repositories/factory.py index 8a37ba933fc..3e8b5277402 100644 --- a/api/core/repositories/factory.py +++ b/api/core/repositories/factory.py @@ -35,6 +35,8 @@ class WorkflowExecutionRepository(Protocol): class WorkflowNodeExecutionRepository(Protocol): def save(self, execution: WorkflowNodeExecution): ... + def save_synchronously(self, execution: WorkflowNodeExecution) -> None: ... + def save_execution_data(self, execution: WorkflowNodeExecution): ... def get_by_workflow_execution( diff --git a/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py b/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py index 15dca1ff8cf..2fe5f427e5d 100644 --- a/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py +++ b/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py @@ -18,6 +18,7 @@ from tenacity import before_sleep_log, retry, retry_if_exception, stop_after_att from configs import dify_config from core.repositories.factory import OrderConfig, WorkflowNodeExecutionRepository +from core.workflow.node_execution_process_data import preserve_workflow_agent_binding_id from extensions.ext_storage import storage from graphon.entities import WorkflowNodeExecution from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus @@ -372,6 +373,12 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) logger.exception("Failed to save workflow node execution after all retries") raise + @override + def save_synchronously(self, execution: WorkflowNodeExecution) -> None: + """Persist a caller row before an Agent v2 participant is materialized.""" + + self.save(execution) + def _persist_to_database(self, db_model: WorkflowNodeExecutionModel): """ Persist the database model to the database. @@ -386,6 +393,13 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) existing = session.get(WorkflowNodeExecutionModel, db_model.id) if existing: + merged_process_data = preserve_workflow_agent_binding_id( + existing.process_data_dict, + db_model.process_data_dict, + ) + db_model.process_data = ( + _deterministic_json_dump(merged_process_data) if merged_process_data is not None else None + ) # Update existing record by copying all non-private attributes for key, value in db_model.__dict__.items(): if not key.startswith("_"): @@ -442,18 +456,25 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) else: db_model.outputs = self._json_encode(domain_model.outputs) - if domain_model.process_data is not None: + process_data = preserve_workflow_agent_binding_id(db_model.process_data_dict, domain_model.process_data) + if process_data is not None: result = self._truncate_and_upload( - domain_model.process_data, + process_data, domain_model.id, ExecutionOffLoadType.PROCESS_DATA, ) if result is not None: - db_model.process_data = self._json_encode(result.truncated_value) - domain_model.set_truncated_process_data(result.truncated_value) + truncated_process_data = preserve_workflow_agent_binding_id( + process_data, + result.truncated_value, + ) + if truncated_process_data is None: + raise ValueError("truncated process data is unavailable") + db_model.process_data = self._json_encode(truncated_process_data) + domain_model.set_truncated_process_data(truncated_process_data) offload_data = _replace_or_append_offload(offload_data, result.offload) else: - db_model.process_data = self._json_encode(domain_model.process_data) + db_model.process_data = self._json_encode(process_data) db_model.offload_data = offload_data with self._session_factory() as session, session.begin(): diff --git a/api/core/workflow/node_execution_process_data.py b/api/core/workflow/node_execution_process_data.py new file mode 100644 index 00000000000..4133feb1b21 --- /dev/null +++ b/api/core/workflow/node_execution_process_data.py @@ -0,0 +1,27 @@ +from collections.abc import Mapping +from typing import Any + +WORKFLOW_AGENT_BINDING_ID_KEY = "workflow_agent_binding_id" + + +def preserve_workflow_agent_binding_id( + identity_source: Mapping[str, Any] | None, + process_data: Mapping[str, Any] | None, +) -> dict[str, Any] | None: + source_id = (identity_source or {}).get(WORKFLOW_AGENT_BINDING_ID_KEY) + target_id = (process_data or {}).get(WORKFLOW_AGENT_BINDING_ID_KEY) + for value in (source_id, target_id): + if value is not None and not isinstance(value, str): + raise ValueError("workflow_agent_binding_id must be a string") + if source_id is not None and target_id is not None and source_id != target_id: + raise ValueError("workflow_agent_binding_id does not match") + + if process_data is None and source_id is None: + return None + merged = dict(process_data or {}) + if source_id is not None: + merged[WORKFLOW_AGENT_BINDING_ID_KEY] = source_id + return merged + + +__all__ = ["WORKFLOW_AGENT_BINDING_ID_KEY", "preserve_workflow_agent_binding_id"] diff --git a/api/core/workflow/node_factory.py b/api/core/workflow/node_factory.py index 02fdb379a45..5b55af4b36b 100644 --- a/api/core/workflow/node_factory.py +++ b/api/core/workflow/node_factory.py @@ -487,7 +487,7 @@ class DifyNodeFactory(NodeFactory): from core.workflow.nodes.agent_v2.output_failure_orchestrator import OutputFailureOrchestrator from core.workflow.nodes.agent_v2.output_file_rebacker import reback_tool_file_output from core.workflow.nodes.agent_v2.output_type_checker import PerOutputTypeChecker - from core.workflow.nodes.agent_v2.session_store import WorkflowAgentRuntimeSessionStore + from core.workflow.nodes.agent_v2.session_store import WorkflowAgentWorkspaceStore return { "binding_resolver": WorkflowAgentBindingResolver(), @@ -512,7 +512,7 @@ class DifyNodeFactory(NodeFactory): # tenant validator resolves ToolFile (canonical) + UploadFile refs. "type_checker": PerOutputTypeChecker(file_validator=AgentOutputFileTenantValidator()), "failure_orchestrator": OutputFailureOrchestrator(), - "session_store": WorkflowAgentRuntimeSessionStore(), + "session_store": WorkflowAgentWorkspaceStore(), } return { "strategy_resolver": self._agent_strategy_resolver, diff --git a/api/core/workflow/nodes/agent_v2/agent_node.py b/api/core/workflow/nodes/agent_v2/agent_node.py index e100b7414b3..9911ef57ad7 100644 --- a/api/core/workflow/nodes/agent_v2/agent_node.py +++ b/api/core/workflow/nodes/agent_v2/agent_node.py @@ -18,13 +18,10 @@ from clients.agent_backend import ( AgentBackendRunEventAdapter, AgentBackendRunFailedInternalEvent, AgentBackendRunSucceededInternalEvent, - AgentBackendSessionCleanupPayload, AgentBackendStreamError, AgentBackendStreamInternalEvent, AgentBackendTransportError, AgentBackendValidationError, - RuntimeLayerSpec, - extract_runtime_layer_specs, ) from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext from core.repositories.human_input_repository import HumanInputFormRepository, HumanInputFormRepositoryImpl @@ -33,11 +30,12 @@ from core.workflow.nodes.human_input.session_binding import default_session_bind from core.workflow.system_variables import SystemVariableKey, get_system_text from graphon.entities.pause_reason import HitlRequired, SchedulingPause from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from graphon.node_events import NodeEventBase, NodeRunResult, PauseRequestedEvent, StreamCompletedEvent +from graphon.graph_events import NodeRunPauseRequestedEvent +from graphon.node_events import NodeEventBase, NodeRunResult, StreamCompletedEvent from graphon.nodes.base.node import Node from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig from services.agent.prompt_mentions import extract_workflow_node_output_selectors -from tasks.agent_backend_session_cleanup_task import cleanup_workflow_agent_runtime_session +from services.agent.workspace_service import AgentWorkspaceNotFoundError from .ask_human_hitl import AskHumanFormBuildError, build_ask_human_pause_reason from .ask_human_resume import build_deferred_tool_results, resolve_ask_human_form @@ -56,7 +54,7 @@ from .runtime_request_builder import ( WorkflowAgentRuntimeRequestBuilder, WorkflowAgentRuntimeRequestBuildError, ) -from .session_store import WorkflowAgentRuntimeSessionStore, WorkflowAgentSessionScope +from .session_store import WorkflowAgentSessionScope, WorkflowAgentWorkspaceStore if TYPE_CHECKING: from graphon.entities import GraphInitParams @@ -68,7 +66,7 @@ logger = logging.getLogger(__name__) # Stage 4 §5+§7: the terminal events that `_consume_event_stream` may return. # Stream + started events are filtered out before we yield; transport errors # are surfaced as a separate StreamCompletedEvent in the second tuple slot. -_TerminalAgentBackendEvent = ( +type _TerminalAgentBackendEvent = ( AgentBackendRunSucceededInternalEvent | AgentBackendRunFailedInternalEvent | AgentBackendRunCancelledInternalEvent @@ -93,7 +91,7 @@ class DifyAgentNode(Node[DifyAgentNodeData]): output_adapter: WorkflowAgentOutputAdapter, type_checker: PerOutputTypeChecker, failure_orchestrator: OutputFailureOrchestrator, - session_store: WorkflowAgentRuntimeSessionStore | None = None, + session_store: WorkflowAgentWorkspaceStore, ) -> None: super().__init__( node_id=node_id, @@ -130,7 +128,34 @@ class DifyAgentNode(Node[DifyAgentNodeData]): return reason @override - def _run(self) -> Generator[NodeEventBase, None, None]: + def _run(self) -> Generator[NodeEventBase | NodeRunPauseRequestedEvent, None, None]: + inputs: dict[str, Any] = {} + process_data: dict[str, Any] = {} + metadata: dict[str, Any] = { + "agent_backend": { + "status": "not_started", + } + } + try: + yield from self._run_inner(inputs=inputs, process_data=process_data, metadata=metadata) + except Exception as error: + if not process_data: + raise + yield self._failure_event( + inputs=inputs, + process_data=process_data, + metadata=metadata, + error=str(error), + error_type="agent_workflow_node_runtime_error", + ) + + def _run_inner( + self, + *, + inputs: dict[str, Any], + process_data: dict[str, Any], + metadata: dict[str, Any], + ) -> Generator[NodeEventBase | NodeRunPauseRequestedEvent, None, None]: dify_ctx = DifyRunContext.model_validate(self.require_run_context_value(DIFY_RUN_CONTEXT_KEY)) workflow_id = self.graph_init_params.workflow_id workflow_run_id = get_system_text( @@ -143,21 +168,24 @@ class DifyAgentNode(Node[DifyAgentNodeData]): self.graph_runtime_state.variable_pool, SystemVariableKey.CONVERSATION_ID, ) - inputs: dict[str, Any] = {} - process_data: dict[str, Any] = {} - metadata: dict[str, Any] = { - "agent_backend": { - "status": "not_started", - } - } # ──── Setup: resolve binding once + extract declared outputs for stage 4 checks ──── try: + existing_scope = self._session_store.load_existing_node_execution_scope( + tenant_id=dify_ctx.tenant_id, + app_id=dify_ctx.app_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run_id, + node_id=self._node_id, + node_execution_id=self.execution_id, + ) bundle = self._binding_resolver.resolve( tenant_id=dify_ctx.tenant_id, app_id=dify_ctx.app_id, workflow_id=workflow_id, node_id=self._node_id, + binding_id=existing_scope.workflow_agent_binding_id if existing_scope is not None else None, + snapshot_id=existing_scope.agent_config_snapshot_id if existing_scope is not None else None, ) except WorkflowAgentBindingError as error: yield self._failure_event( @@ -168,20 +196,31 @@ class DifyAgentNode(Node[DifyAgentNodeData]): error_type=error.error_code, ) return + except AgentWorkspaceNotFoundError as error: + yield self._failure_event( + inputs=inputs, + process_data=process_data, + metadata=metadata, + error=str(error), + error_type="agent_workflow_node_runtime_error", + ) + return - process_data = { - "agent_id": bundle.agent.id, - "agent_config_snapshot_id": bundle.snapshot.id, - "binding_id": bundle.binding.id, - } - session_scope = WorkflowAgentSessionScope( + process_data.update( + { + "agent_id": bundle.agent.id, + "agent_config_snapshot_id": bundle.snapshot.id, + "workflow_agent_binding_id": bundle.binding.id, + } + ) + session_scope = existing_scope or WorkflowAgentSessionScope( tenant_id=dify_ctx.tenant_id, app_id=dify_ctx.app_id, workflow_id=workflow_id, workflow_run_id=workflow_run_id, node_id=self._node_id, - node_execution_id=self.id, - binding_id=bundle.binding.id, + node_execution_id=self.execution_id, + workflow_agent_binding_id=bundle.binding.id, agent_id=bundle.agent.id, agent_config_snapshot_id=bundle.snapshot.id, ) @@ -200,47 +239,53 @@ class DifyAgentNode(Node[DifyAgentNodeData]): # the second Agent run as deferred_tool_results; if it is somehow still # waiting, re-emit the same pause defensively. deferred_tool_results = None - if self._session_store is not None: - stored_session = self._session_store.load_active_session(session_scope) - if stored_session is not None and stored_session.pending_form_id is not None: - resume_outcome = resolve_ask_human_form( - form_id=stored_session.pending_form_id, - tenant_id=dify_ctx.tenant_id, - node_id=self._node_id, + stored_session = self._session_store.load_or_create_node_execution_session( + session_scope, + home_snapshot_id=bundle.snapshot.home_snapshot_id, + ) + if stored_session.pending_form_id is not None: + resume_outcome = resolve_ask_human_form( + form_id=stored_session.pending_form_id, + tenant_id=dify_ctx.tenant_id, + node_id=self._node_id, + ) + if resume_outcome is not None and resume_outcome.repause is not None: + yield self._pause_event( + reason=resume_outcome.repause, + inputs=inputs, + process_data=process_data, + metadata=metadata, + ) + return + if ( + resume_outcome is not None + and resume_outcome.deferred_result is not None + and stored_session.pending_tool_call_id is not None + ): + deferred_tool_results = build_deferred_tool_results( + tool_call_id=stored_session.pending_tool_call_id, + result=resume_outcome.deferred_result, ) - if resume_outcome is not None and resume_outcome.repause is not None: - yield PauseRequestedEvent(reason=self._to_graph_pause_reason(resume_outcome.repause)) - return - if ( - resume_outcome is not None - and resume_outcome.deferred_result is not None - and stored_session.pending_tool_call_id is not None - ): - deferred_tool_results = build_deferred_tool_results( - tool_call_id=stored_session.pending_tool_call_id, - result=resume_outcome.deferred_result, - ) # ──── Retry loop (Stage 4 §7) ──── attempt = 0 while True: try: - session_snapshot = None - if self._session_store is not None: - session_snapshot = self._session_store.load_active_snapshot(session_scope) runtime_request = self._runtime_request_builder.build( WorkflowAgentRuntimeBuildContext( dify_context=dify_ctx, workflow_id=workflow_id, workflow_run_id=workflow_run_id, node_id=self._node_id, - node_execution_id=self.id, + node_execution_id=self.execution_id, variable_pool=self.graph_runtime_state.variable_pool, binding=bundle.binding, agent=bundle.agent, snapshot=bundle.snapshot, + binding_id=stored_session.binding_id, + backend_binding_ref=stored_session.backend_binding_ref, attempt=attempt, - session_snapshot=session_snapshot, + session_snapshot=stored_session.session_snapshot, deferred_tool_results=deferred_tool_results, ) ) @@ -266,8 +311,9 @@ class DifyAgentNode(Node[DifyAgentNodeData]): # Capture inputs only from the first attempt so retry doesn't churn the # node's "inputs" payload that ends up in the workflow detail view. if attempt == 0: - inputs = {"agent_backend_request": runtime_request.redacted_request} - metadata = dict(runtime_request.metadata) + inputs["agent_backend_request"] = runtime_request.redacted_request + metadata.clear() + metadata.update(runtime_request.metadata) metadata["attempt"] = attempt try: @@ -288,7 +334,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]): "status": create_response.status, } - terminal_event, exhausted = self._consume_event_stream(create_response.run_id, metadata) + terminal_event, exhausted = self._consume_event_stream( + create_response.run_id, + inputs=inputs, + process_data=process_data, + metadata=metadata, + ) if exhausted is not None: # Streaming error / unexpected end — surface immediately without # retrying because the failure is transport-level. @@ -349,29 +400,23 @@ class DifyAgentNode(Node[DifyAgentNodeData]): ) self._save_session_snapshot( session_scope=session_scope, - backend_run_id=terminal_event.run_id, + binding_id=stored_session.binding_id, snapshot=terminal_event.session_snapshot, - runtime_layer_specs=extract_runtime_layer_specs(runtime_request.request.composition), metadata=metadata, pending_form_id=pending_form_id, pending_tool_call_id=pending_tool_call_id, ) - yield PauseRequestedEvent(reason=self._to_graph_pause_reason(pause_reason)) - return - - # Non-success terminal (failed / cancelled) skips per-output - # post-processing — the backend itself already failed. We also retire - # the local ACTIVE session row so a workflow loop back into the same - # Agent node cannot resume from a stale snapshot. The failed agent - # backend layers (suspended per ``on_exit``) are left for agent - # backend's own GC; this row will no longer be picked up by the - # workflow-terminal cleanup layer. - if not isinstance(terminal_event, AgentBackendRunSucceededInternalEvent): - self._mark_session_cleaned_on_failure( - session_scope=session_scope, - backend_run_id=terminal_event.run_id, + yield self._pause_event( + reason=pause_reason, + inputs=inputs, + process_data=process_data, metadata=metadata, ) + return + + # A failed attempt does not retire the product-owned Binding. The + # Workflow Run terminal lifecycle event owns that transition. + if not isinstance(terminal_event, AgentBackendRunSucceededInternalEvent): yield StreamCompletedEvent( node_run_result=self._output_adapter.build_failure_result( event=terminal_event, @@ -384,9 +429,8 @@ class DifyAgentNode(Node[DifyAgentNodeData]): self._save_session_snapshot( session_scope=session_scope, - backend_run_id=terminal_event.run_id, + binding_id=stored_session.binding_id, snapshot=terminal_event.session_snapshot, - runtime_layer_specs=extract_runtime_layer_specs(runtime_request.request.composition), metadata=metadata, ) @@ -458,6 +502,9 @@ class DifyAgentNode(Node[DifyAgentNodeData]): def _consume_event_stream( self, run_id: str, + *, + inputs: dict[str, Any], + process_data: dict[str, Any], metadata: dict[str, Any], ) -> tuple[ _TerminalAgentBackendEvent | None, @@ -507,8 +554,8 @@ class DifyAgentNode(Node[DifyAgentNodeData]): return internal_event, None self._cancel_backend_run(run_id, reason="unexpected_event") return None, self._failure_event( - inputs={}, - process_data={}, + inputs=inputs, + process_data=process_data, metadata=metadata, error=f"Unexpected internal event type {internal_event.type!r}", error_type="agent_backend_stream_error", @@ -516,8 +563,8 @@ class DifyAgentNode(Node[DifyAgentNodeData]): except AgentBackendError as error: self._cancel_backend_run(run_id, reason=self._stream_stop_reason()) return None, self._failure_event( - inputs={}, - process_data={}, + inputs=inputs, + process_data=process_data, metadata=metadata, error=str(error), error_type=self._agent_backend_error_type(error), @@ -525,8 +572,8 @@ class DifyAgentNode(Node[DifyAgentNodeData]): except Exception as error: self._cancel_backend_run(run_id, reason=self._stream_stop_reason()) return None, self._failure_event( - inputs={}, - process_data={}, + inputs=inputs, + process_data=process_data, metadata=metadata, error=str(error), error_type="agent_backend_stream_error", @@ -596,21 +643,17 @@ class DifyAgentNode(Node[DifyAgentNodeData]): self, *, session_scope: WorkflowAgentSessionScope, - backend_run_id: str, + binding_id: str, snapshot: CompositorSessionSnapshot | None, - runtime_layer_specs: list[RuntimeLayerSpec], metadata: dict[str, Any], pending_form_id: str | None = None, pending_tool_call_id: str | None = None, ) -> None: - if self._session_store is None: - return try: self._session_store.save_active_snapshot( scope=session_scope, - backend_run_id=backend_run_id, + binding_id=binding_id, snapshot=snapshot, - runtime_layer_specs=runtime_layer_specs, pending_form_id=pending_form_id, pending_tool_call_id=pending_tool_call_id, ) @@ -619,88 +662,18 @@ class DifyAgentNode(Node[DifyAgentNodeData]): metadata["agent_backend"] = agent_backend except Exception: logger.warning( - "Failed to persist workflow Agent runtime session snapshot: " - "tenant_id=%s workflow_run_id=%s node_id=%s binding_id=%s agent_id=%s backend_run_id=%s", + "Failed to persist workflow Agent Binding session snapshot: " + "tenant_id=%s workflow_run_id=%s node_id=%s binding_id=%s agent_id=%s", session_scope.tenant_id, session_scope.workflow_run_id, session_scope.node_id, - session_scope.binding_id, + session_scope.workflow_agent_binding_id, session_scope.agent_id, - backend_run_id, exc_info=True, ) agent_backend = dict(metadata.get("agent_backend") or {}) agent_backend["session_snapshot_persisted"] = False - agent_backend["session_snapshot_persist_error"] = "workflow_agent_runtime_session_store_error" - metadata["agent_backend"] = agent_backend - - def _mark_session_cleaned_on_failure( - self, - *, - session_scope: WorkflowAgentSessionScope, - backend_run_id: str, - metadata: dict[str, Any], - ) -> None: - if self._session_store is None: - return - stored_session = self._session_store.load_active_session(session_scope) - try: - if stored_session is not None and stored_session.runtime_layer_specs: - payload = AgentBackendSessionCleanupPayload( - session_snapshot=stored_session.session_snapshot, - runtime_layer_specs=stored_session.runtime_layer_specs, - idempotency_key=( - f"{session_scope.tenant_id}:{session_scope.workflow_run_id}:{session_scope.node_id}:" - f"{session_scope.binding_id}:workflow-agent-failure-cleanup:" - f"{stored_session.backend_run_id or 'no-stored-run'}:{backend_run_id}" - ), - metadata={ - "tenant_id": session_scope.tenant_id, - "app_id": session_scope.app_id, - "workflow_id": session_scope.workflow_id, - "workflow_run_id": session_scope.workflow_run_id, - "node_id": session_scope.node_id, - "node_execution_id": session_scope.node_execution_id, - "binding_id": session_scope.binding_id, - "agent_id": session_scope.agent_id, - "agent_config_snapshot_id": session_scope.agent_config_snapshot_id, - "previous_agent_backend_run_id": stored_session.backend_run_id, - "failed_agent_backend_run_id": backend_run_id, - }, - ) - cleanup_workflow_agent_runtime_session.delay(payload.model_dump(mode="json")) - except Exception: - logger.warning( - "Failed to enqueue workflow Agent backend cleanup on agent run failure: " - "tenant_id=%s workflow_run_id=%s node_id=%s binding_id=%s agent_id=%s backend_run_id=%s", - session_scope.tenant_id, - session_scope.workflow_run_id, - session_scope.node_id, - session_scope.binding_id, - session_scope.agent_id, - backend_run_id, - exc_info=True, - ) - try: - self._session_store.mark_cleaned(scope=session_scope, backend_run_id=backend_run_id) - agent_backend = dict(metadata.get("agent_backend") or {}) - agent_backend["session_snapshot_cleaned_on_failure"] = True - metadata["agent_backend"] = agent_backend - except Exception: - logger.warning( - "Failed to mark workflow Agent runtime session cleaned on agent run failure: " - "tenant_id=%s workflow_run_id=%s node_id=%s binding_id=%s agent_id=%s backend_run_id=%s", - session_scope.tenant_id, - session_scope.workflow_run_id, - session_scope.node_id, - session_scope.binding_id, - session_scope.agent_id, - backend_run_id, - exc_info=True, - ) - agent_backend = dict(metadata.get("agent_backend") or {}) - agent_backend["session_snapshot_cleaned_on_failure"] = False - agent_backend["session_snapshot_cleanup_error"] = "workflow_agent_runtime_session_store_error" + agent_backend["session_snapshot_persist_error"] = "workflow_agent_workspace_store_error" metadata["agent_backend"] = agent_backend @staticmethod @@ -742,6 +715,27 @@ class DifyAgentNode(Node[DifyAgentNodeData]): ) ) + def _pause_event( + self, + *, + reason: HumanInputRequired | SchedulingPause, + inputs: dict[str, Any], + process_data: dict[str, Any], + metadata: dict[str, Any], + ) -> NodeRunPauseRequestedEvent: + return NodeRunPauseRequestedEvent( + id=self.execution_id, + node_id=self._node_id, + node_type=self.node_type, + node_run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.PAUSED, + inputs=inputs, + process_data=process_data, + metadata={WorkflowNodeExecutionMetadataKey.AGENT_LOG: metadata}, + ), + reason=self._to_graph_pause_reason(reason), + ) + @staticmethod def _agent_backend_error_type(error: AgentBackendError) -> str: if isinstance(error, AgentBackendValidationError): diff --git a/api/core/workflow/nodes/agent_v2/binding_resolver.py b/api/core/workflow/nodes/agent_v2/binding_resolver.py index 34812f670f7..3710af100b0 100644 --- a/api/core/workflow/nodes/agent_v2/binding_resolver.py +++ b/api/core/workflow/nodes/agent_v2/binding_resolver.py @@ -41,18 +41,27 @@ class WorkflowAgentBindingResolver: app_id: str, workflow_id: str, node_id: str, + binding_id: str | None = None, + snapshot_id: str | None = None, ) -> WorkflowAgentBindingBundle: - with session_factory.create_session() as session: - binding = session.scalar( - select(WorkflowAgentNodeBinding) - .where( - WorkflowAgentNodeBinding.tenant_id == tenant_id, - WorkflowAgentNodeBinding.app_id == app_id, - WorkflowAgentNodeBinding.workflow_id == workflow_id, - WorkflowAgentNodeBinding.node_id == node_id, - ) - .limit(1) + """Resolve the current binding, optionally at a generation pinned by an existing execution.""" + + if (binding_id is None) != (snapshot_id is None): + raise WorkflowAgentBindingError( + "agent_binding_generation_invalid", + "Workflow Agent binding and config snapshot must be pinned together.", ) + + with session_factory.create_session() as session: + binding_stmt = select(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.tenant_id == tenant_id, + WorkflowAgentNodeBinding.app_id == app_id, + WorkflowAgentNodeBinding.workflow_id == workflow_id, + WorkflowAgentNodeBinding.node_id == node_id, + ) + if binding_id is not None: + binding_stmt = binding_stmt.where(WorkflowAgentNodeBinding.id == binding_id) + binding = session.scalar(binding_stmt.limit(1)) if binding is None: raise WorkflowAgentBindingError( "agent_binding_not_found", @@ -77,12 +86,16 @@ class WorkflowAgentBindingResolver: f"Agent {binding.agent_id} is not available or has not been published.", ) - snapshot_id = ( - agent.active_config_snapshot_id - if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT - else binding.current_snapshot_id + effective_snapshot_id = ( + ( + agent.active_config_snapshot_id + if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT + else binding.current_snapshot_id + ) + if snapshot_id is None + else snapshot_id ) - if snapshot_id is None: + if effective_snapshot_id is None: raise WorkflowAgentBindingError( "agent_config_snapshot_not_found", "Workflow Agent binding has no current config snapshot.", @@ -93,14 +106,14 @@ class WorkflowAgentBindingResolver: .where( AgentConfigSnapshot.tenant_id == tenant_id, AgentConfigSnapshot.agent_id == agent.id, - AgentConfigSnapshot.id == snapshot_id, + AgentConfigSnapshot.id == effective_snapshot_id, ) .limit(1) ) if snapshot is None: raise WorkflowAgentBindingError( "agent_config_snapshot_not_found", - f"Agent config snapshot {snapshot_id} not found.", + f"Agent config snapshot {effective_snapshot_id} not found.", ) session.expunge(binding) diff --git a/api/core/workflow/nodes/agent_v2/runtime_request_builder.py b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py index 22b99412053..09860998c89 100644 --- a/api/core/workflow/nodes/agent_v2/runtime_request_builder.py +++ b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py @@ -33,7 +33,6 @@ from dify_agent.layers.shell import ( DifyShellCliToolConfig, DifyShellEnvVarConfig, DifyShellLayerConfig, - DifyShellSandboxConfig, DifyShellSecretRefConfig, ) from dify_agent.protocol import CreateRunRequest, DeferredToolResultsPayload @@ -136,6 +135,8 @@ class WorkflowAgentRuntimeBuildContext: binding: WorkflowAgentNodeBinding agent: Agent snapshot: AgentConfigSnapshot + binding_id: str + backend_binding_ref: str # Stage 4 §7 / D-4: 0 for the first run, then incremented per retry. Drives the # idempotency key so the backend treats each retry as a fresh request. attempt: int = 0 @@ -251,6 +252,7 @@ class WorkflowAgentRuntimeRequestBuilder: agent_mode=self._agent_backend_agent_mode(context.dify_context.invoke_from), invoke_from=cast(DifyExecutionContextInvokeFrom, context.dify_context.invoke_from.value), ), + backend_binding_ref=context.backend_binding_ref, agent_soul_prompt=soul_prompt or None, workflow_node_job_prompt=workflow_job_prompt, user_prompt=user_prompt, @@ -739,7 +741,6 @@ class WorkflowAgentRuntimeRequestBuilder: def build_shell_layer_config(agent_soul: AgentSoulConfig) -> DifyShellLayerConfig: """Map Agent Soul shell-adjacent fields into the Agent backend shell config.""" - sandbox_config = _plain_mapping(agent_soul.sandbox.config) return DifyShellLayerConfig( cli_tools=[ tool @@ -750,12 +751,6 @@ def build_shell_layer_config(agent_soul: AgentSoulConfig) -> DifyShellLayerConfi secret_refs=[ secret for secret in (_shell_secret_ref(item) for item in agent_soul.env.secret_refs) if secret is not None ], - sandbox=DifyShellSandboxConfig( - provider=agent_soul.sandbox.provider, - config=sandbox_config, - ) - if agent_soul.sandbox.provider or sandbox_config - else None, ) diff --git a/api/core/workflow/nodes/agent_v2/session_cleanup_layer.py b/api/core/workflow/nodes/agent_v2/session_cleanup_layer.py deleted file mode 100644 index c3a7c23a56c..00000000000 --- a/api/core/workflow/nodes/agent_v2/session_cleanup_layer.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Workflow terminal layer that retires Agent backend sessions asynchronously.""" - -from __future__ import annotations - -import logging -from typing import override - -from clients.agent_backend import AgentBackendSessionCleanupPayload -from core.workflow.system_variables import SystemVariableKey, get_system_text -from graphon.graph_engine.layers import GraphEngineLayer -from graphon.graph_events import ( - GraphEngineEvent, - GraphRunAbortedEvent, - GraphRunFailedEvent, - GraphRunPartialSucceededEvent, - GraphRunSucceededEvent, -) -from tasks.agent_backend_session_cleanup_task import cleanup_workflow_agent_runtime_session - -from .session_store import StoredWorkflowAgentSession, WorkflowAgentRuntimeSessionStore - -logger = logging.getLogger(__name__) - - -class WorkflowAgentSessionCleanupLayer(GraphEngineLayer): - """Retire workflow-owned Agent runtime sessions when the workflow ends. - - Workflow termination is a product-lifecycle boundary: once the run reaches a - terminal graph event, the local session row must no longer be resumable. The - actual Agent backend cleanup is therefore dispatched asynchronously with the - persisted snapshot/specs payload, while the local row is marked CLEANED - immediately afterwards regardless of enqueue outcome. - """ - - _TERMINAL_EVENTS = ( - GraphRunSucceededEvent, - GraphRunPartialSucceededEvent, - GraphRunFailedEvent, - GraphRunAbortedEvent, - ) - - def __init__(self, *, session_store: WorkflowAgentRuntimeSessionStore) -> None: - super().__init__() - self._session_store = session_store - - @override - def on_graph_start(self) -> None: - return - - @override - def on_event(self, event: GraphEngineEvent) -> None: - if not isinstance(event, self._TERMINAL_EVENTS): - return - workflow_run_id = get_system_text( - self.graph_runtime_state.variable_pool, - SystemVariableKey.WORKFLOW_EXECUTION_ID, - ) - if not workflow_run_id: - logger.warning("Skipping workflow Agent session cleanup: workflow_run_id is missing.") - return - - for stored_session in self._session_store.list_active_sessions(workflow_run_id=workflow_run_id): - self._cleanup_session(stored_session) - - @override - def on_graph_end(self, error: Exception | None) -> None: - return - - def _cleanup_session(self, stored_session: StoredWorkflowAgentSession) -> None: - scope = stored_session.scope - try: - if stored_session.runtime_layer_specs: - payload = AgentBackendSessionCleanupPayload( - session_snapshot=stored_session.session_snapshot, - runtime_layer_specs=stored_session.runtime_layer_specs, - idempotency_key=f"{scope.workflow_run_id}:{scope.node_id}:{scope.binding_id}:agent-session-cleanup", - metadata={ - "tenant_id": scope.tenant_id, - "app_id": scope.app_id, - "workflow_id": scope.workflow_id, - "workflow_run_id": scope.workflow_run_id, - "node_id": scope.node_id, - "node_execution_id": scope.node_execution_id, - "binding_id": scope.binding_id, - "agent_id": scope.agent_id, - "agent_config_snapshot_id": scope.agent_config_snapshot_id, - "previous_agent_backend_run_id": stored_session.backend_run_id, - }, - ) - cleanup_workflow_agent_runtime_session.delay(payload.model_dump(mode="json")) - else: - logger.warning( - "Skipping workflow Agent backend cleanup enqueue: no runtime_layer_specs persisted. " - "workflow_run_id=%s node_id=%s agent_id=%s", - scope.workflow_run_id, - scope.node_id, - scope.agent_id, - ) - except Exception: - logger.warning( - "Failed to enqueue workflow Agent backend cleanup: " - "workflow_run_id=%s node_id=%s agent_id=%s previous_run_id=%s", - scope.workflow_run_id, - scope.node_id, - scope.agent_id, - stored_session.backend_run_id, - exc_info=True, - ) - finally: - try: - self._session_store.mark_cleaned(scope=scope, backend_run_id=stored_session.backend_run_id) - except Exception: - logger.warning( - "Failed to retire workflow Agent runtime session after cleanup enqueue: " - "workflow_run_id=%s node_id=%s agent_id=%s previous_run_id=%s", - scope.workflow_run_id, - scope.node_id, - scope.agent_id, - stored_session.backend_run_id, - exc_info=True, - ) - - -def build_workflow_agent_session_cleanup_layer() -> WorkflowAgentSessionCleanupLayer: - """Wire the cleanup layer with the standard workflow-owned session store.""" - return WorkflowAgentSessionCleanupLayer(session_store=WorkflowAgentRuntimeSessionStore()) diff --git a/api/core/workflow/nodes/agent_v2/session_store.py b/api/core/workflow/nodes/agent_v2/session_store.py index 215fa67b5eb..2952d5abc60 100644 --- a/api/core/workflow/nodes/agent_v2/session_store.py +++ b/api/core/workflow/nodes/agent_v2/session_store.py @@ -1,31 +1,32 @@ +"""Workflow Agent participant persistence keyed by node execution.""" + from __future__ import annotations -from dataclasses import dataclass, field +import json +import time +from dataclasses import dataclass from agenton.compositor import CompositorSessionSnapshot -from dify_agent.protocol import RuntimeLayerSpec -from pydantic import TypeAdapter from sqlalchemy import select +from sqlalchemy.orm import Session from core.db.session_factory import session_factory -from libs.datetime_utils import naive_utc_now from models.agent import ( - AgentRuntimeSessionOwnerType, - WorkflowAgentRuntimeSession, - WorkflowAgentRuntimeSessionStatus, + AgentConfigVersionKind, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) +from models.workflow import WorkflowNodeExecutionModel +from services.agent.workspace_service import ( + AgentWorkspaceNotFoundError, + AgentWorkspaceService, + WorkspaceOwnerScope, ) -_SPECS_ADAPTER: TypeAdapter[list[RuntimeLayerSpec]] = TypeAdapter(list[RuntimeLayerSpec]) - - -def _serialize_specs(specs: list[RuntimeLayerSpec]) -> str: - return _SPECS_ADAPTER.dump_json(specs).decode() - - -def _deserialize_specs(value: str | None) -> list[RuntimeLayerSpec]: - if not value: - return [] - return _SPECS_ADAPTER.validate_json(value) +_CALLER_VISIBILITY_ATTEMPTS = 60 +_CALLER_VISIBILITY_INTERVAL_SECONDS = 0.05 @dataclass(frozen=True, slots=True) @@ -36,171 +37,251 @@ class WorkflowAgentSessionScope: workflow_run_id: str | None node_id: str node_execution_id: str - binding_id: str + workflow_agent_binding_id: str agent_id: str agent_config_snapshot_id: str + @property + def workspace_owner(self) -> WorkspaceOwnerScope: + return WorkspaceOwnerScope( + tenant_id=self.tenant_id, + app_id=self.app_id, + owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN, + owner_id=self.workflow_run_id or self.node_execution_id, + owner_scope_key=f"{self.node_id}:{self.workflow_agent_binding_id}", + ) + @dataclass(frozen=True, slots=True) class StoredWorkflowAgentSession: scope: WorkflowAgentSessionScope - session_snapshot: CompositorSessionSnapshot - backend_run_id: str | None - runtime_layer_specs: list[RuntimeLayerSpec] = field(default_factory=list) - # ENG-637: set while the session is paused on a dify.ask_human deferred call. + binding_id: str + workspace_id: str + backend_binding_ref: str + session_snapshot: CompositorSessionSnapshot | None pending_form_id: str | None = None pending_tool_call_id: str | None = None -class WorkflowAgentRuntimeSessionStore: - """Stores Agent backend session snapshots for workflow Agent node re-entry.""" +class WorkflowAgentWorkspaceStore: + """Load or create the participant named by a node execution caller row.""" - def load_active_snapshot(self, scope: WorkflowAgentSessionScope) -> CompositorSessionSnapshot | None: - stored = self.load_active_session(scope) - return stored.session_snapshot if stored is not None else None - - def load_active_session(self, scope: WorkflowAgentSessionScope) -> StoredWorkflowAgentSession | None: - """Load the active session row including any pending ask_human correlation.""" - if scope.workflow_run_id is None: - return None + def load_existing_node_execution_scope( + self, + *, + tenant_id: str, + app_id: str, + workflow_id: str, + workflow_run_id: str | None, + node_id: str, + node_execution_id: str, + ) -> WorkflowAgentSessionScope | None: + """Return the generation pinned by an existing node execution participant.""" with session_factory.create_session() as session: - row = session.scalar( - select(WorkflowAgentRuntimeSession).where( - WorkflowAgentRuntimeSession.tenant_id == scope.tenant_id, - WorkflowAgentRuntimeSession.workflow_run_id == scope.workflow_run_id, - WorkflowAgentRuntimeSession.node_id == scope.node_id, - WorkflowAgentRuntimeSession.binding_id == scope.binding_id, - WorkflowAgentRuntimeSession.agent_id == scope.agent_id, - WorkflowAgentRuntimeSession.status == WorkflowAgentRuntimeSessionStatus.ACTIVE, - ) + execution = self._load_execution_by_identity( + session=session, + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run_id, + node_id=node_id, + node_execution_id=node_execution_id, ) - if row is None: + binding_id = execution.agent_workspace_binding_id + if binding_id is None: return None - return StoredWorkflowAgentSession( - scope=scope, - session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), - backend_run_id=row.backend_run_id, - runtime_layer_specs=_deserialize_specs(row.composition_layer_specs), - pending_form_id=row.pending_form_id, - pending_tool_call_id=row.pending_tool_call_id, + process_data = execution.process_data_dict + if not isinstance(process_data, dict): + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is invalid") + workflow_agent_binding_id = process_data.get("workflow_agent_binding_id") + if not isinstance(workflow_agent_binding_id, str): + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is missing") + owner_scope = WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN, + owner_id=workflow_run_id or node_execution_id, + owner_scope_key=f"{node_id}:{workflow_agent_binding_id}", + ) + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=tenant_id, + binding_id=binding_id, + expected_owner_scope=owner_scope, + ) + if binding is None or binding.agent_config_version_kind != AgentConfigVersionKind.SNAPSHOT: + raise AgentWorkspaceNotFoundError("Workflow node participant Binding is unavailable") + return WorkflowAgentSessionScope( + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run_id, + node_id=node_id, + node_execution_id=node_execution_id, + workflow_agent_binding_id=workflow_agent_binding_id, + agent_id=binding.agent_id, + agent_config_snapshot_id=binding.agent_config_version_id, ) - def list_active_sessions(self, *, workflow_run_id: str) -> list[StoredWorkflowAgentSession]: + def load_or_create_node_execution_session( + self, scope: WorkflowAgentSessionScope, *, home_snapshot_id: str + ) -> StoredWorkflowAgentSession: with session_factory.create_session() as session: - rows = session.scalars( - select(WorkflowAgentRuntimeSession).where( - WorkflowAgentRuntimeSession.workflow_run_id == workflow_run_id, - WorkflowAgentRuntimeSession.status == WorkflowAgentRuntimeSessionStatus.ACTIVE, + execution = self._load_execution(session=session, scope=scope) + process_data = execution.process_data_dict + if process_data is None: + process_data = {} + if not isinstance(process_data, dict): + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is invalid") + stored_workflow_binding_id = process_data.get("workflow_agent_binding_id") + if stored_workflow_binding_id is not None and stored_workflow_binding_id != scope.workflow_agent_binding_id: + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity does not match") + + binding_id = execution.agent_workspace_binding_id + if binding_id is None: + binding = AgentWorkspaceService.create_binding( + session=session, + scope=scope.workspace_owner, + agent_id=scope.agent_id, + base_home_snapshot_id=home_snapshot_id, + agent_config_version_id=scope.agent_config_snapshot_id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, ) - ).all() - return [ - StoredWorkflowAgentSession( - scope=WorkflowAgentSessionScope( - tenant_id=row.tenant_id, - app_id=row.app_id, - # These columns are nullable on the unified runtime-session - # table (workflow_run ⊕ conversation owner), but are always - # populated for a workflow-owned row; coerce for the typed scope. - workflow_id=row.workflow_id or "", - workflow_run_id=row.workflow_run_id, - node_id=row.node_id or "", - node_execution_id=row.node_execution_id or "", - binding_id=row.binding_id or "", - agent_id=row.agent_id, - agent_config_snapshot_id=row.agent_config_snapshot_id or "", - ), - session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), - backend_run_id=row.backend_run_id, - runtime_layer_specs=_deserialize_specs(row.composition_layer_specs), + execution.agent_workspace_binding_id = binding.id + execution.process_data = json.dumps( + { + **process_data, + "workflow_agent_binding_id": scope.workflow_agent_binding_id, + }, + ensure_ascii=False, ) - for row in rows - ] + session.commit() + else: + if stored_workflow_binding_id is None: + raise AgentWorkspaceNotFoundError("Workflow node execution caller identity is missing") + resolved_binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=scope.tenant_id, + binding_id=binding_id, + expected_owner_scope=scope.workspace_owner, + ) + if resolved_binding is None or resolved_binding.agent_id != scope.agent_id: + raise AgentWorkspaceNotFoundError("Workflow node participant Binding is unavailable") + binding = resolved_binding + AgentWorkspaceService.validate_binding_generation( + binding, + base_home_snapshot_id=home_snapshot_id, + agent_config_version_id=scope.agent_config_snapshot_id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + return self._stored(scope, binding) def save_active_snapshot( self, *, scope: WorkflowAgentSessionScope, - backend_run_id: str, + binding_id: str, snapshot: CompositorSessionSnapshot | None, - runtime_layer_specs: list[RuntimeLayerSpec], pending_form_id: str | None = None, pending_tool_call_id: str | None = None, ) -> None: - if scope.workflow_run_id is None or snapshot is None: + if snapshot is None: return + AgentWorkspaceService.save_binding_session_snapshot( + tenant_id=scope.tenant_id, + binding_id=binding_id, + session_snapshot=snapshot.model_dump_json(), + pending_form_id=pending_form_id, + pending_tool_call_id=pending_tool_call_id, + ) - snapshot_json = snapshot.model_dump_json() - specs_json = _serialize_specs(runtime_layer_specs) + def retire_workflow_run(self, *, tenant_id: str, app_id: str, workflow_run_id: str) -> list[str]: + """Retire active Workspaces, commit, and return active or already-retired IDs for collection.""" + + retired: list[str] = [] with session_factory.create_session() as session: - row = session.scalar( - select(WorkflowAgentRuntimeSession).where( - WorkflowAgentRuntimeSession.tenant_id == scope.tenant_id, - WorkflowAgentRuntimeSession.workflow_run_id == scope.workflow_run_id, - WorkflowAgentRuntimeSession.node_id == scope.node_id, - WorkflowAgentRuntimeSession.binding_id == scope.binding_id, - WorkflowAgentRuntimeSession.agent_id == scope.agent_id, + workspaces = session.scalars( + select(AgentWorkspace).where( + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.app_id == app_id, + AgentWorkspace.owner_type == AgentWorkspaceOwnerType.WORKFLOW_RUN, + AgentWorkspace.owner_id == workflow_run_id, + AgentWorkspace.status.in_((AgentWorkingResourceStatus.ACTIVE, AgentWorkingResourceStatus.RETIRED)), ) - ) - if row is None: - row = WorkflowAgentRuntimeSession( - tenant_id=scope.tenant_id, - app_id=scope.app_id, - owner_type=AgentRuntimeSessionOwnerType.WORKFLOW_RUN, - workflow_id=scope.workflow_id, - workflow_run_id=scope.workflow_run_id, - node_id=scope.node_id, - node_execution_id=scope.node_execution_id, - binding_id=scope.binding_id, - agent_id=scope.agent_id, - agent_config_snapshot_id=scope.agent_config_snapshot_id, - backend_run_id=backend_run_id, - session_snapshot=snapshot_json, - composition_layer_specs=specs_json, - status=WorkflowAgentRuntimeSessionStatus.ACTIVE, - pending_form_id=pending_form_id, - pending_tool_call_id=pending_tool_call_id, + ).all() + for workspace in workspaces: + if workspace.status == AgentWorkingResourceStatus.RETIRED: + retired.append(workspace.id) + continue + workspace_id = AgentWorkspaceService.retire_workspace( + session=session, + tenant_id=tenant_id, + workspace_id=workspace.id, ) - session.add(row) - else: - row.node_execution_id = scope.node_execution_id - row.agent_config_snapshot_id = scope.agent_config_snapshot_id - row.backend_run_id = backend_run_id - row.session_snapshot = snapshot_json - row.composition_layer_specs = specs_json - row.status = WorkflowAgentRuntimeSessionStatus.ACTIVE - row.cleaned_at = None - # Set (or clear, when omitted) the ask_human pause correlation. - row.pending_form_id = pending_form_id - row.pending_tool_call_id = pending_tool_call_id + if workspace_id is not None: + retired.append(workspace_id) session.commit() + return retired - def mark_cleaned(self, *, scope: WorkflowAgentSessionScope, backend_run_id: str | None = None) -> None: - if scope.workflow_run_id is None: - return + @staticmethod + def _load_execution(*, session: Session, scope: WorkflowAgentSessionScope) -> WorkflowNodeExecutionModel: + return WorkflowAgentWorkspaceStore._load_execution_by_identity( + session=session, + tenant_id=scope.tenant_id, + app_id=scope.app_id, + workflow_id=scope.workflow_id, + workflow_run_id=scope.workflow_run_id, + node_id=scope.node_id, + node_execution_id=scope.node_execution_id, + ) - with session_factory.create_session() as session: - row = session.scalar( - select(WorkflowAgentRuntimeSession).where( - WorkflowAgentRuntimeSession.tenant_id == scope.tenant_id, - WorkflowAgentRuntimeSession.workflow_run_id == scope.workflow_run_id, - WorkflowAgentRuntimeSession.node_id == scope.node_id, - WorkflowAgentRuntimeSession.binding_id == scope.binding_id, - WorkflowAgentRuntimeSession.agent_id == scope.agent_id, - WorkflowAgentRuntimeSession.status == WorkflowAgentRuntimeSessionStatus.ACTIVE, - ) - ) - if row is None: - return - if backend_run_id is not None: - row.backend_run_id = backend_run_id - row.status = WorkflowAgentRuntimeSessionStatus.CLEANED - row.cleaned_at = naive_utc_now() - session.commit() + @staticmethod + def _load_execution_by_identity( + *, + session: Session, + tenant_id: str, + app_id: str, + workflow_id: str, + workflow_run_id: str | None, + node_id: str, + node_execution_id: str, + ) -> WorkflowNodeExecutionModel: + """Wait briefly for the already-emitted node-start event to persist its caller row.""" + + stmt = select(WorkflowNodeExecutionModel).where( + WorkflowNodeExecutionModel.id == node_execution_id, + WorkflowNodeExecutionModel.tenant_id == tenant_id, + WorkflowNodeExecutionModel.app_id == app_id, + WorkflowNodeExecutionModel.workflow_id == workflow_id, + WorkflowNodeExecutionModel.node_id == node_id, + WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id, + ) + for attempt in range(_CALLER_VISIBILITY_ATTEMPTS): + execution = session.scalar(stmt) + if execution is not None: + return execution + if attempt < _CALLER_VISIBILITY_ATTEMPTS - 1: + time.sleep(_CALLER_VISIBILITY_INTERVAL_SECONDS) + + raise AgentWorkspaceNotFoundError("Workflow node execution caller is unavailable") + + @staticmethod + def _stored(scope: WorkflowAgentSessionScope, binding: AgentWorkspaceBinding) -> StoredWorkflowAgentSession: + snapshot = ( + CompositorSessionSnapshot.model_validate_json(binding.session_snapshot) + if binding.session_snapshot + else None + ) + return StoredWorkflowAgentSession( + scope=scope, + binding_id=binding.id, + workspace_id=binding.workspace_id, + backend_binding_ref=binding.backend_binding_ref, + session_snapshot=snapshot, + pending_form_id=binding.pending_form_id, + pending_tool_call_id=binding.pending_tool_call_id, + ) -__all__ = [ - "StoredWorkflowAgentSession", - "WorkflowAgentRuntimeSessionStore", - "WorkflowAgentSessionScope", -] +__all__ = ["StoredWorkflowAgentSession", "WorkflowAgentSessionScope", "WorkflowAgentWorkspaceStore"] diff --git a/api/core/workflow/nodes/agent_v2/workspace_retirement_layer.py b/api/core/workflow/nodes/agent_v2/workspace_retirement_layer.py new file mode 100644 index 00000000000..7c44897fcb7 --- /dev/null +++ b/api/core/workflow/nodes/agent_v2/workspace_retirement_layer.py @@ -0,0 +1,89 @@ +"""Retire Workflow Agent Workspaces when the Workflow Run terminates.""" + +from __future__ import annotations + +import logging +from typing import override + +from core.app.entities.app_invoke_entities import DifyRunContext +from core.workflow.nodes.agent_v2.session_store import WorkflowAgentWorkspaceStore +from core.workflow.system_variables import SystemVariableKey, get_system_text +from graphon.graph_engine.layers import GraphEngineLayer +from graphon.graph_events import ( + GraphEngineEvent, + GraphRunAbortedEvent, + GraphRunFailedEvent, + GraphRunPartialSucceededEvent, + GraphRunSucceededEvent, +) +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection + +logger = logging.getLogger(__name__) + + +class WorkflowAgentWorkspaceRetirementLayer(GraphEngineLayer): + """Synchronously retire run Workspaces, then enqueue physical collection.""" + + _TERMINAL_EVENTS = ( + GraphRunSucceededEvent, + GraphRunPartialSucceededEvent, + GraphRunFailedEvent, + GraphRunAbortedEvent, + ) + + def __init__( + self, + *, + dify_run_context: DifyRunContext, + ) -> None: + super().__init__() + self._dify_run_context = dify_run_context + + @override + def on_graph_start(self) -> None: + return + + @override + def on_event(self, event: GraphEngineEvent) -> None: + if not isinstance(event, self._TERMINAL_EVENTS): + return + workflow_run_id = get_system_text( + self.graph_runtime_state.variable_pool, + SystemVariableKey.WORKFLOW_EXECUTION_ID, + ) + if not workflow_run_id: + logger.warning("Skipping Workflow Agent Workspace retirement: workflow_run_id is missing") + return + try: + workspace_ids = WorkflowAgentWorkspaceStore().retire_workflow_run( + tenant_id=self._dify_run_context.tenant_id, + app_id=self._dify_run_context.app_id, + workflow_run_id=workflow_run_id, + ) + except Exception: + logger.exception( + "Failed to retire Workflow Agent Workspaces", + extra={ + "tenant_id": self._dify_run_context.tenant_id, + "app_id": self._dify_run_context.app_id, + "workflow_run_id": workflow_run_id, + }, + ) + return + enqueue_agent_resource_collection( + tenant_id=self._dify_run_context.tenant_id, + workspace_ids=workspace_ids, + ) + + @override + def on_graph_end(self, error: Exception | None) -> None: + return + + +def build_workflow_agent_workspace_retirement_layer( + *, dify_run_context: DifyRunContext +) -> WorkflowAgentWorkspaceRetirementLayer: + return WorkflowAgentWorkspaceRetirementLayer(dify_run_context=dify_run_context) + + +__all__ = ["WorkflowAgentWorkspaceRetirementLayer", "build_workflow_agent_workspace_retirement_layer"] diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 690fa64cdf8..a96d50e1caf 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -166,6 +166,7 @@ def init_app(app: DifyApp) -> Celery: imports = [ "tasks.async_workflow_tasks", # trigger workers + "tasks.collect_agent_resources_task", # retired Agent resource collection "tasks.trigger_processing_tasks", # async trigger processing "tasks.generate_summary_index_task", # summary index generation "tasks.regenerate_summary_index_task", # summary index regeneration diff --git a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py index 1c62df3746d..fa11e210c05 100644 --- a/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py +++ b/api/extensions/logstore/repositories/logstore_workflow_node_execution_repository.py @@ -277,6 +277,12 @@ class LogstoreWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): logger.exception("Failed to dual-write node execution to SQL database: id=%s", execution.id) # Don't raise - LogStore write succeeded, SQL is just a backup + @override + def save_synchronously(self, execution: WorkflowNodeExecution) -> None: + """Create the SQL caller row required by Agent v2 participant ownership.""" + + self.sql_repository.save_synchronously(execution) + @override def save_execution_data(self, execution: WorkflowNodeExecution) -> None: """ diff --git a/api/migrations/versions/2026_07_21_2251-2f39536b3feb_add_agent_home_snapshot_ledger.py b/api/migrations/versions/2026_07_21_2251-2f39536b3feb_add_agent_home_snapshot_ledger.py new file mode 100644 index 00000000000..70fc241b5b2 --- /dev/null +++ b/api/migrations/versions/2026_07_21_2251-2f39536b3feb_add_agent_home_snapshot_ledger.py @@ -0,0 +1,64 @@ +"""add agent home snapshot ledger + +Revision ID: 2f39536b3feb +Revises: 6f5a9c2d8e1b +Create Date: 2026-07-21 22:51:07.268658 + +""" +from alembic import op +import models as models +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '2f39536b3feb' +down_revision = '6f5a9c2d8e1b' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('agent_home_snapshots', + sa.Column('id', models.types.StringUUID(), nullable=False), + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('agent_id', models.types.StringUUID(), nullable=False), + sa.Column('snapshot_ref', sa.String(length=255), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.PrimaryKeyConstraint('id', name='agent_home_snapshot_pkey') + ) + with op.batch_alter_table('agent_home_snapshots', schema=None) as batch_op: + batch_op.create_index('agent_home_snapshot_tenant_agent_idx', ['tenant_id', 'agent_id'], unique=False) + + with op.batch_alter_table('agent_config_drafts', schema=None) as batch_op: + batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=False)) + + with op.batch_alter_table('agent_config_snapshots', schema=None) as batch_op: + batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=False)) + + with op.batch_alter_table('agent_runtime_sessions', schema=None) as batch_op: + batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=False)) + batch_op.drop_index(batch_op.f('agent_runtime_session_conversation_scope_unique'), postgresql_where='(conversation_id IS NOT NULL)') + batch_op.create_index('agent_runtime_session_conversation_scope_unique', ['tenant_id', 'conversation_id', 'agent_id', 'agent_config_snapshot_id', 'home_snapshot_id'], unique=True, postgresql_where=sa.text('conversation_id IS NOT NULL')) + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('agent_runtime_sessions', schema=None) as batch_op: + batch_op.drop_index('agent_runtime_session_conversation_scope_unique', postgresql_where=sa.text('conversation_id IS NOT NULL')) + batch_op.create_index(batch_op.f('agent_runtime_session_conversation_scope_unique'), ['tenant_id', 'conversation_id', 'agent_id', 'agent_config_snapshot_id'], unique=True, postgresql_where='(conversation_id IS NOT NULL)') + batch_op.drop_column('home_snapshot_id') + + with op.batch_alter_table('agent_config_snapshots', schema=None) as batch_op: + batch_op.drop_column('home_snapshot_id') + + with op.batch_alter_table('agent_config_drafts', schema=None) as batch_op: + batch_op.drop_column('home_snapshot_id') + + with op.batch_alter_table('agent_home_snapshots', schema=None) as batch_op: + batch_op.drop_index('agent_home_snapshot_tenant_agent_idx') + + op.drop_table('agent_home_snapshots') + # ### end Alembic commands ### diff --git a/api/migrations/versions/2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py b/api/migrations/versions/2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py new file mode 100644 index 00000000000..7a7ec50a696 --- /dev/null +++ b/api/migrations/versions/2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py @@ -0,0 +1,153 @@ +"""replace agent runtime sessions with workspaces and bindings + +Revision ID: f6e4c5686857 +Revises: 2f39536b3feb +Create Date: 2026-07-23 02:03:05.641638 + +""" +from alembic import op +import models as models +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = 'f6e4c5686857' +down_revision = '2f39536b3feb' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('agent_workspace_bindings', + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('app_id', models.types.StringUUID(), nullable=False), + sa.Column('workspace_id', models.types.StringUUID(), nullable=False), + sa.Column('agent_id', models.types.StringUUID(), nullable=False), + sa.Column('base_home_snapshot_id', models.types.StringUUID(), nullable=False), + sa.Column('agent_config_version_id', models.types.StringUUID(), nullable=False), + sa.Column('agent_config_version_kind', sa.String(length=32), nullable=False), + sa.Column('backend_binding_ref', sa.String(length=255), nullable=False), + sa.Column('session_snapshot', models.types.LongText(), nullable=True), + sa.Column('status', sa.String(length=32), server_default='active', nullable=False), + sa.Column('retired_at', sa.DateTime(), nullable=True), + sa.Column('pending_form_id', models.types.StringUUID(), nullable=True), + sa.Column('pending_tool_call_id', sa.String(length=255), nullable=True), + sa.Column('id', models.types.StringUUID(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.PrimaryKeyConstraint('id', name='agent_workspace_binding_pkey') + ) + with op.batch_alter_table('agent_workspace_bindings', schema=None) as batch_op: + batch_op.create_index('agent_workspace_binding_agent_status_idx', ['tenant_id', 'agent_id', 'status'], unique=False) + batch_op.create_index('agent_workspace_binding_status_retired_idx', ['status', 'retired_at'], unique=False) + batch_op.create_index('agent_workspace_binding_workspace_status_idx', ['tenant_id', 'workspace_id', 'status'], unique=False) + + op.create_table('agent_workspaces', + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('app_id', models.types.StringUUID(), nullable=False), + sa.Column('owner_type', sa.String(length=32), nullable=False), + sa.Column('owner_id', models.types.StringUUID(), nullable=False), + sa.Column('owner_scope_key', sa.String(length=255), nullable=False), + sa.Column('backend_workspace_ref', sa.String(length=255), nullable=False), + sa.Column('status', sa.String(length=32), server_default='active', nullable=False), + sa.Column('active_guard', sa.SmallInteger(), server_default='1', nullable=True), + sa.Column('retired_at', sa.DateTime(), nullable=True), + sa.Column('id', models.types.StringUUID(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.PrimaryKeyConstraint('id', name='agent_workspace_pkey') + ) + with op.batch_alter_table('agent_workspaces', schema=None) as batch_op: + batch_op.create_index('agent_workspace_owner_active_unique', ['tenant_id', 'owner_type', 'owner_id', 'owner_scope_key', 'active_guard'], unique=True) + batch_op.create_index('agent_workspace_status_retired_idx', ['status', 'retired_at'], unique=False) + batch_op.create_index('agent_workspace_tenant_app_status_idx', ['tenant_id', 'app_id', 'status'], unique=False) + batch_op.create_index('agent_workspace_tenant_status_idx', ['tenant_id', 'status'], unique=False) + + with op.batch_alter_table('agent_runtime_sessions', schema=None) as batch_op: + batch_op.drop_index(batch_op.f('agent_runtime_session_backend_run_idx')) + batch_op.drop_index(batch_op.f('agent_runtime_session_conversation_lookup_idx')) + batch_op.drop_index(batch_op.f('agent_runtime_session_conversation_scope_unique'), postgresql_where='(conversation_id IS NOT NULL)') + batch_op.drop_index(batch_op.f('agent_runtime_session_workflow_lookup_idx')) + batch_op.drop_index(batch_op.f('agent_runtime_session_workflow_scope_unique'), postgresql_where='(workflow_run_id IS NOT NULL)') + + op.drop_table('agent_runtime_sessions') + with op.batch_alter_table('agent_home_snapshots', schema=None) as batch_op: + batch_op.add_column(sa.Column('status', sa.String(length=32), server_default='active', nullable=False)) + batch_op.add_column(sa.Column('retired_at', sa.DateTime(), nullable=True)) + batch_op.create_index('agent_home_snapshot_status_retired_idx', ['status', 'retired_at'], unique=False) + + with op.batch_alter_table('conversations', schema=None) as batch_op: + batch_op.add_column(sa.Column('agent_workspace_binding_id', models.types.StringUUID(), nullable=True)) + + with op.batch_alter_table('agent_config_drafts', schema=None) as batch_op: + batch_op.add_column(sa.Column('agent_workspace_binding_id', models.types.StringUUID(), nullable=True)) + + with op.batch_alter_table('workflow_node_executions', schema=None) as batch_op: + batch_op.add_column(sa.Column('agent_workspace_binding_id', models.types.StringUUID(), nullable=True)) + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('workflow_node_executions', schema=None) as batch_op: + batch_op.drop_column('agent_workspace_binding_id') + + with op.batch_alter_table('agent_config_drafts', schema=None) as batch_op: + batch_op.drop_column('agent_workspace_binding_id') + + with op.batch_alter_table('conversations', schema=None) as batch_op: + batch_op.drop_column('agent_workspace_binding_id') + + with op.batch_alter_table('agent_home_snapshots', schema=None) as batch_op: + batch_op.drop_index('agent_home_snapshot_status_retired_idx') + batch_op.drop_column('retired_at') + batch_op.drop_column('status') + + op.create_table('agent_runtime_sessions', + sa.Column('id', sa.UUID(), server_default=sa.text('uuidv7()'), autoincrement=False, nullable=False), + sa.Column('tenant_id', sa.UUID(), autoincrement=False, nullable=False), + sa.Column('app_id', sa.UUID(), autoincrement=False, nullable=False), + sa.Column('owner_type', sa.VARCHAR(length=32), autoincrement=False, nullable=False), + sa.Column('agent_id', sa.UUID(), autoincrement=False, nullable=False), + sa.Column('backend_run_id', sa.VARCHAR(length=255), autoincrement=False, nullable=True), + sa.Column('session_snapshot', sa.TEXT(), autoincrement=False, nullable=False), + sa.Column('workflow_id', sa.UUID(), autoincrement=False, nullable=True), + sa.Column('workflow_run_id', sa.UUID(), autoincrement=False, nullable=True), + sa.Column('node_id', sa.VARCHAR(length=255), autoincrement=False, nullable=True), + sa.Column('node_execution_id', sa.VARCHAR(length=255), autoincrement=False, nullable=True), + sa.Column('binding_id', sa.UUID(), autoincrement=False, nullable=True), + sa.Column('agent_config_snapshot_id', sa.UUID(), autoincrement=False, nullable=True), + sa.Column('composition_layer_specs', sa.TEXT(), autoincrement=False, nullable=False), + sa.Column('conversation_id', sa.UUID(), autoincrement=False, nullable=True), + sa.Column('status', sa.VARCHAR(length=32), server_default=sa.text("'active'::character varying"), autoincrement=False, nullable=False), + sa.Column('cleaned_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True), + sa.Column('created_at', postgresql.TIMESTAMP(), server_default=sa.text('CURRENT_TIMESTAMP'), autoincrement=False, nullable=False), + sa.Column('updated_at', postgresql.TIMESTAMP(), server_default=sa.text('CURRENT_TIMESTAMP'), autoincrement=False, nullable=False), + sa.Column('pending_form_id', sa.UUID(), autoincrement=False, nullable=True), + sa.Column('pending_tool_call_id', sa.VARCHAR(length=255), autoincrement=False, nullable=True), + sa.Column('home_snapshot_id', sa.UUID(), autoincrement=False, nullable=False), + sa.PrimaryKeyConstraint('id', name=op.f('agent_runtime_session_pkey')) + ) + with op.batch_alter_table('agent_runtime_sessions', schema=None) as batch_op: + batch_op.create_index(batch_op.f('agent_runtime_session_workflow_scope_unique'), ['tenant_id', 'workflow_run_id', 'node_id', 'binding_id', 'agent_id'], unique=True, postgresql_where='(workflow_run_id IS NOT NULL)') + batch_op.create_index(batch_op.f('agent_runtime_session_workflow_lookup_idx'), ['tenant_id', 'workflow_run_id', 'node_id', 'status'], unique=False) + batch_op.create_index(batch_op.f('agent_runtime_session_conversation_scope_unique'), ['tenant_id', 'conversation_id', 'agent_id', 'agent_config_snapshot_id', 'home_snapshot_id'], unique=True, postgresql_where='(conversation_id IS NOT NULL)') + batch_op.create_index(batch_op.f('agent_runtime_session_conversation_lookup_idx'), ['tenant_id', 'conversation_id', 'status'], unique=False) + batch_op.create_index(batch_op.f('agent_runtime_session_backend_run_idx'), ['backend_run_id'], unique=False) + + with op.batch_alter_table('agent_workspaces', schema=None) as batch_op: + batch_op.drop_index('agent_workspace_tenant_status_idx') + batch_op.drop_index('agent_workspace_tenant_app_status_idx') + batch_op.drop_index('agent_workspace_status_retired_idx') + batch_op.drop_index('agent_workspace_owner_active_unique') + + op.drop_table('agent_workspaces') + with op.batch_alter_table('agent_workspace_bindings', schema=None) as batch_op: + batch_op.drop_index('agent_workspace_binding_workspace_status_idx') + batch_op.drop_index('agent_workspace_binding_status_retired_idx') + batch_op.drop_index('agent_workspace_binding_agent_status_idx') + + op.drop_table('agent_workspace_bindings') + # ### end Alembic commands ### diff --git a/api/models/__init__.py b/api/models/__init__.py index b4c1362b414..387a6761820 100644 --- a/api/models/__init__.py +++ b/api/models/__init__.py @@ -15,21 +15,22 @@ from .agent import ( AgentConfigRevision, AgentConfigRevisionOperation, AgentConfigSnapshot, + AgentConfigVersionKind, AgentDebugConversation, AgentDriveFile, AgentDriveFileKind, + AgentHomeSnapshot, AgentIconType, AgentKind, - AgentRuntimeSession, - AgentRuntimeSessionOwnerType, - AgentRuntimeSessionStatus, AgentScope, AgentSource, AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, WorkflowAgentBindingType, WorkflowAgentNodeBinding, - WorkflowAgentRuntimeSession, - WorkflowAgentRuntimeSessionStatus, ) from .api_based_extension import APIBasedExtension, APIBasedExtensionPoint from .comment import ( @@ -164,17 +165,20 @@ __all__ = [ "AgentConfigRevision", "AgentConfigRevisionOperation", "AgentConfigSnapshot", + "AgentConfigVersionKind", "AgentDebugConversation", "AgentDriveFile", "AgentDriveFileKind", + "AgentHomeSnapshot", "AgentIconType", "AgentKind", - "AgentRuntimeSession", - "AgentRuntimeSessionOwnerType", - "AgentRuntimeSessionStatus", "AgentScope", "AgentSource", "AgentStatus", + "AgentWorkingResourceStatus", + "AgentWorkspace", + "AgentWorkspaceBinding", + "AgentWorkspaceOwnerType", "ApiRequest", "ApiToken", "ApiToolProvider", @@ -271,8 +275,6 @@ __all__ = [ "Workflow", "WorkflowAgentBindingType", "WorkflowAgentNodeBinding", - "WorkflowAgentRuntimeSession", - "WorkflowAgentRuntimeSessionStatus", "WorkflowAppLog", "WorkflowAppLogCreatedFrom", "WorkflowArchiveLog", diff --git a/api/models/agent.py b/api/models/agent.py index cd3d371481d..2ecfcaf896d 100644 --- a/api/models/agent.py +++ b/api/models/agent.py @@ -116,35 +116,29 @@ class WorkflowAgentBindingType(StrEnum): INLINE_AGENT = "inline_agent" -class AgentRuntimeSessionStatus(StrEnum): - """Lifecycle state of an Agent backend session snapshot. +class AgentWorkingResourceStatus(StrEnum): + """Product lifecycle state for a persistent working-environment resource.""" - Owner-agnostic: applies both to workflow Agent Node runs (owner = - workflow_run) and to Agent App conversations (owner = conversation). - """ - - # Snapshot can be reused by a later Agent run in the same session. ACTIVE = "active" - # Snapshot has been retired and must not be submitted to Agent backend again. - CLEANED = "cleaned" + RETIRED = "retired" -class AgentRuntimeSessionOwnerType(StrEnum): - """Which product surface owns an Agent runtime session row.""" +class AgentWorkspaceOwnerType(StrEnum): + """Product scope that owns a Workspace.""" - # Owned by one workflow Agent Node execution scope. WORKFLOW_RUN = "workflow_run" - # Owned by one Agent App conversation (multi-turn chat). CONVERSATION = "conversation" + BUILD_DRAFT = "build_draft" -# Back-compat alias: the workflow lifecycle code (shipped in PR #36724) imports -# the old name. Kept so unifying the table does not churn that path. -WorkflowAgentRuntimeSessionStatus = AgentRuntimeSessionStatus +class AgentConfigVersionKind(StrEnum): + SNAPSHOT = "snapshot" + DRAFT = "draft" + BUILD_DRAFT = "build_draft" class Agent(DefaultFieldsMixin, Base): - """Workspace-scoped Agent identity used by Agent Roster and workflow-only agents.""" + """Agent Soul and source lineage; ``AgentWorkspaceBinding.id`` identifies each materialized participant.""" __tablename__ = "agents" __table_args__ = ( @@ -221,14 +215,42 @@ class Agent(DefaultFieldsMixin, Base): archived_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) -class AgentDebugConversation(DefaultFieldsMixin, Base): - """Per-account, per-draft console debug conversation for an Agent App. +class AgentHomeSnapshot(Base): + """Append-only mapping from one Agent-owned Home identity to its backend ref. - Agent App preview state must be isolated by editor account. The Agent row is - shared by everyone in the workspace, so this table owns the user-specific - conversation pointers used by console debug chat. ``draft`` is the Preview - conversation and ``debug_build`` is the Build conversation; they must never - share persisted messages or runtime sessions. + Product tables reference ``id``. ``snapshot_ref`` remains an opaque + deployment-specific handle and is only consumed at Dify Agent boundaries. + Snapshot bytes and ``snapshot_ref`` are immutable. Lifecycle metadata can + transition ACTIVE -> RETIRED; successful physical collection deletes row. + """ + + __tablename__ = "agent_home_snapshots" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="agent_home_snapshot_pkey"), + Index("agent_home_snapshot_tenant_agent_idx", "tenant_id", "agent_id"), + Index("agent_home_snapshot_status_retired_idx", "status", "retired_at"), + ) + + id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7())) + tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + agent_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + snapshot_ref: Mapped[str] = mapped_column(String(255), nullable=False) + status: Mapped[AgentWorkingResourceStatus] = mapped_column( + EnumText(AgentWorkingResourceStatus, length=32), + nullable=False, + default=AgentWorkingResourceStatus.ACTIVE, + server_default=AgentWorkingResourceStatus.ACTIVE.value, + ) + retired_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) + created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp()) + + +class AgentDebugConversation(DefaultFieldsMixin, Base): + """Current console Conversation pointer for one account and draft surface. + + This row owns no Binding or runtime. A Preview Conversation holds its + CONVERSATION Binding pointer, while a DEBUG_BUILD AgentConfigDraft holds its + BUILD_DRAFT Binding pointer. """ __tablename__ = "agent_debug_conversations" @@ -259,7 +281,11 @@ class AgentDebugConversation(DefaultFieldsMixin, Base): class AgentConfigDraft(DefaultFieldsMixin, Base): - """Editable Agent Soul draft separated from immutable published snapshots.""" + """Editable Agent Soul draft separated from immutable published snapshots. + + A DEBUG_BUILD draft owns its materialized participant through + ``agent_workspace_binding_id``. Normal drafts leave that pointer unset. + """ __tablename__ = "agent_config_drafts" __table_args__ = ( @@ -281,6 +307,8 @@ class AgentConfigDraft(DefaultFieldsMixin, Base): account_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) draft_owner_key: Mapped[str] = mapped_column(String(255), nullable=False, default="") base_snapshot_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) + home_snapshot_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + agent_workspace_binding_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) config_snapshot: Mapped[Any] = mapped_column(JSONModelColumn(AgentSoulConfig), nullable=False) created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True) updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True) @@ -316,6 +344,7 @@ class AgentConfigSnapshot(DefaultFieldsMixin, Base): agent_id: Mapped[str] = mapped_column(StringUUID, nullable=False) version: Mapped[int] = mapped_column(sa.Integer, nullable=False) config_snapshot: Mapped[Any] = mapped_column(JSONModelColumn(AgentSoulConfig), nullable=False) + home_snapshot_id: Mapped[str] = mapped_column(StringUUID, nullable=False) summary: Mapped[str | None] = mapped_column(LongText, nullable=True) version_note: Mapped[str | None] = mapped_column(LongText, nullable=True) created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True) @@ -432,102 +461,83 @@ class WorkflowAgentNodeBinding(DefaultFieldsMixin, Base): return dict(self.node_job_config) -class AgentRuntimeSession(DefaultFieldsMixin, Base): - """Persisted Agent backend session snapshot, owner-agnostic. +class AgentWorkspace(DefaultFieldsMixin, Base): + """Mutable Workspace owned by one product scope, independent of Agents.""" - One unified table serves both owners (decision Q2): - - workflow Agent Node runs: ``owner_type = workflow_run``; the - ``workflow_id / workflow_run_id / node_id / binding_id / - agent_config_snapshot_id / composition_layer_specs`` columns are set. - - Agent App conversations: ``owner_type = conversation``; the - ``conversation_id`` column is set and the workflow columns stay NULL. - Runtime state is scoped by ``agent_config_snapshot_id``. For published - web/API runs this points to an immutable AgentConfigSnapshot; for console - debugger/build runs it points to the editable AgentConfigDraft row. - - The snapshot is runtime state returned by Agent backend, kept separate from - Agent Soul snapshots and workflow node-job config. - """ - - __tablename__ = "agent_runtime_sessions" + __tablename__ = "agent_workspaces" __table_args__ = ( - sa.PrimaryKeyConstraint("id", name="agent_runtime_session_pkey"), - # Workflow owner uniqueness (partial: only rows with a workflow_run_id). + sa.PrimaryKeyConstraint("id", name="agent_workspace_pkey"), Index( - "agent_runtime_session_workflow_scope_unique", + "agent_workspace_owner_active_unique", "tenant_id", - "workflow_run_id", - "node_id", - "binding_id", - "agent_id", + "owner_type", + "owner_id", + "owner_scope_key", + "active_guard", unique=True, - postgresql_where=sa.text("workflow_run_id IS NOT NULL"), ), - # Conversation owner uniqueness (partial: only rows with a conversation_id). - Index( - "agent_runtime_session_conversation_scope_unique", - "tenant_id", - "conversation_id", - "agent_id", - "agent_config_snapshot_id", - unique=True, - postgresql_where=sa.text("conversation_id IS NOT NULL"), - ), - Index( - "agent_runtime_session_workflow_lookup_idx", - "tenant_id", - "workflow_run_id", - "node_id", - "status", - ), - Index( - "agent_runtime_session_conversation_lookup_idx", - "tenant_id", - "conversation_id", - "status", - ), - Index("agent_runtime_session_backend_run_idx", "backend_run_id"), + Index("agent_workspace_tenant_status_idx", "tenant_id", "status"), + Index("agent_workspace_tenant_app_status_idx", "tenant_id", "app_id", "status"), + Index("agent_workspace_status_retired_idx", "status", "retired_at"), ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) app_id: Mapped[str] = mapped_column(StringUUID, nullable=False) - owner_type: Mapped[AgentRuntimeSessionOwnerType] = mapped_column( - EnumText(AgentRuntimeSessionOwnerType, length=32), nullable=False + owner_type: Mapped[AgentWorkspaceOwnerType] = mapped_column( + EnumText(AgentWorkspaceOwnerType, length=32), nullable=False ) - agent_id: Mapped[str] = mapped_column(StringUUID, nullable=False) - backend_run_id: Mapped[str | None] = mapped_column(String(255), nullable=True) - session_snapshot: Mapped[str] = mapped_column(LongText, nullable=False) - # Workflow-owner columns (NULL for conversation owner). - workflow_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) - workflow_run_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) - node_id: Mapped[str | None] = mapped_column(String(255), nullable=True) - node_execution_id: Mapped[str | None] = mapped_column(String(255), nullable=True) - binding_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) - agent_config_snapshot_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) - # JSON-encoded list of non-sensitive runtime layer specs ({name, type, deps, - # config}). The persisted schema keeps its original name because the sandbox - # refactor intentionally avoids a storage migration. - composition_layer_specs: Mapped[str] = mapped_column(LongText, nullable=False, server_default="[]") - # Conversation-owner column (NULL for workflow owner). - conversation_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) - status: Mapped[AgentRuntimeSessionStatus] = mapped_column( - EnumText(AgentRuntimeSessionStatus, length=32), + owner_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + owner_scope_key: Mapped[str] = mapped_column(String(255), nullable=False) + backend_workspace_ref: Mapped[str] = mapped_column(String(255), nullable=False) + status: Mapped[AgentWorkingResourceStatus] = mapped_column( + EnumText(AgentWorkingResourceStatus, length=32), nullable=False, - default=AgentRuntimeSessionStatus.ACTIVE, + default=AgentWorkingResourceStatus.ACTIVE, + server_default=AgentWorkingResourceStatus.ACTIVE.value, ) - cleaned_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) - # ENG-637: when a run pauses for a dify.ask_human deferred call, these link - # the session to the awaiting HITL form and the deferred tool_call_id, so a - # resumed node can map the submitted form back into deferred_tool_results. - # Both NULL whenever the session is not paused on human input. + active_guard: Mapped[int | None] = mapped_column(sa.SmallInteger, nullable=True, default=1, server_default="1") + retired_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) + + +class AgentWorkspaceBinding(DefaultFieldsMixin, Base): + """One materialized Agent participant and session attached to a Workspace. + + All resource IDs are logical associations rather than database foreign + keys, so RETIRED rows can outlive their Workspace or base Home Snapshot. + ``agent_id`` identifies the source Agent Soul; this row's ``id`` identifies + the participant and its private Materialized Home. + """ + + __tablename__ = "agent_workspace_bindings" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="agent_workspace_binding_pkey"), + Index("agent_workspace_binding_workspace_status_idx", "tenant_id", "workspace_id", "status"), + Index("agent_workspace_binding_agent_status_idx", "tenant_id", "agent_id", "status"), + Index("agent_workspace_binding_status_retired_idx", "status", "retired_at"), + ) + + tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + app_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + workspace_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + agent_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + base_home_snapshot_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + agent_config_version_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + agent_config_version_kind: Mapped[AgentConfigVersionKind] = mapped_column( + EnumText(AgentConfigVersionKind, length=32), nullable=False + ) + backend_binding_ref: Mapped[str] = mapped_column(String(255), nullable=False) + session_snapshot: Mapped[str | None] = mapped_column(LongText, nullable=True) + status: Mapped[AgentWorkingResourceStatus] = mapped_column( + EnumText(AgentWorkingResourceStatus, length=32), + nullable=False, + default=AgentWorkingResourceStatus.ACTIVE, + server_default=AgentWorkingResourceStatus.ACTIVE.value, + ) + retired_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) pending_form_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) pending_tool_call_id: Mapped[str | None] = mapped_column(String(255), nullable=True) -# Back-compat alias for the shipped workflow lifecycle code (PR #36724). -WorkflowAgentRuntimeSession = AgentRuntimeSession - - class AgentDriveFileKind(StrEnum): """Kind of existing file record an agent-drive KV entry points at.""" diff --git a/api/models/model.py b/api/models/model.py index 07ac06284cb..73006d1fb58 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -1160,6 +1160,13 @@ class OAuthProviderApp(TypeBase): class Conversation(Base): + """Conversation state, including the exact Agent participant when applicable. + + ``agent_workspace_binding_id`` is a logical pointer rather than a foreign + key because retired Binding ledger rows may be collected before the + conversation history is deleted. + """ + __tablename__ = "conversations" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="conversation_pkey"), @@ -1181,6 +1188,7 @@ class Conversation(Base): id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4())) app_id = mapped_column(StringUUID, nullable=False) app_model_config_id = mapped_column(StringUUID, nullable=True) + agent_workspace_binding_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) model_provider = mapped_column(String(255), nullable=True) override_model_configs = mapped_column(LongText) model_id = mapped_column(String(255), nullable=True) diff --git a/api/models/workflow.py b/api/models/workflow.py index 58135e79334..b902a3978ef 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -1031,6 +1031,7 @@ class WorkflowNodeExecutionModel(Base): # This model is expected to have `offlo node_id: Mapped[str] = mapped_column(String(255)) node_type: Mapped[str] = mapped_column(String(255)) title: Mapped[str] = mapped_column(String(255)) + agent_workspace_binding_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) inputs: Mapped[str | None] = mapped_column(LongText) process_data: Mapped[str | None] = mapped_column(LongText) outputs: Mapped[str | None] = mapped_column(LongText) diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index e4d46101bb3..4a3e277f4fb 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -964,12 +964,6 @@ Stop a running Agent App chat message generation | ---- | ---------- | ----------- | -------- | ------ | | agent_id | path | | Yes | string (uuid) | -#### Request Body - -| Required | Schema | -| -------- | ------ | -| No | **application/json**: [AgentDebugConversationRefreshPayload](#agentdebugconversationrefreshpayload)
| - #### Responses | Code | Description | Schema | @@ -1261,7 +1255,8 @@ Get basic information for an Agent App conversation sandbox | Name | Located in | Description | Required | Schema | | ---- | ---------- | ----------- | -------- | ------ | | agent_id | path | Agent ID | Yes | string (uuid) | -| conversation_id | query | Agent App conversation ID | Yes | string | +| caller_id | query | Agent App caller ID | Yes | string | +| caller_type | query | | Yes | string,
**Available values:** "build_draft", "conversation" | #### Responses @@ -1277,7 +1272,8 @@ List a directory in an Agent App conversation sandbox | Name | Located in | Description | Required | Schema | | ---- | ---------- | ----------- | -------- | ------ | | agent_id | path | Agent ID | Yes | string (uuid) | -| conversation_id | query | Agent App conversation ID | Yes | string | +| caller_id | query | Agent App caller ID | Yes | string | +| caller_type | query | | Yes | string,
**Available values:** "build_draft", "conversation" | | path | query | Directory path relative to the sandbox workspace | No | string,
**Default:** . | #### Responses @@ -1294,7 +1290,8 @@ Read a text/binary preview file in an Agent App conversation sandbox | Name | Located in | Description | Required | Schema | | ---- | ---------- | ----------- | -------- | ------ | | agent_id | path | Agent ID | Yes | string (uuid) | -| conversation_id | query | Agent App conversation ID | Yes | string | +| caller_id | query | Agent App caller ID | Yes | string | +| caller_type | query | | Yes | string,
**Available values:** "build_draft", "conversation" | | path | query | File path relative to the sandbox workspace | Yes | string | #### Responses @@ -3807,7 +3804,7 @@ List a directory in a workflow Agent node sandbox | app_id | path | Application ID | Yes | string (uuid) | | node_id | path | Workflow Agent node ID | Yes | string | | workflow_run_id | path | Workflow run ID | Yes | string (uuid) | -| node_execution_id | query | Optional workflow node execution ID. When omitted, the latest active session for the node is used. | No | string | +| node_execution_id | query | Workflow node execution ID | Yes | string | | path | query | Directory path relative to the sandbox workspace | No | string,
**Default:** . | #### Responses @@ -3826,7 +3823,7 @@ Read a text/binary preview file in a workflow Agent node sandbox | app_id | path | Application ID | Yes | string (uuid) | | node_id | path | Workflow Agent node ID | Yes | string | | workflow_run_id | path | Workflow run ID | Yes | string (uuid) | -| node_execution_id | query | Optional workflow node execution ID. When omitted, the latest active session for the node is used. | No | string | +| node_execution_id | query | Workflow node execution ID | Yes | string | | path | query | File path relative to the sandbox workspace | Yes | string | #### Responses @@ -13943,12 +13940,6 @@ Stable Agent Soul reference to one normalized skill archive. | date | string | | Yes | | message_count | integer | | Yes | -#### AgentDebugConversationRefreshPayload - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| draft_type | [AgentConfigDraftType](#agentconfigdrafttype) | Agent draft surface whose conversation should be refreshed | No | - #### AgentDebugConversationRefreshResponse | Name | Type | Description | Required | @@ -14666,7 +14657,8 @@ section may be empty, which is how callers express "no knowledge layer". | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| conversation_id | string | Agent App conversation ID | Yes | +| caller_id | string | Agent App caller ID | Yes | +| caller_type | string,
**Available values:** "build_draft", "conversation" | *Enum:* `"build_draft"`, `"conversation"` | Yes | | path | string | File path relative to the sandbox workspace | Yes | #### AgentScope @@ -21347,7 +21339,6 @@ Whitelist scopes accepted by RBAC app and dataset access config APIs. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| session_id | string | | Yes | | workspace_cwd | string | | Yes | #### SandboxListResponse @@ -23167,7 +23158,7 @@ How a workflow node is bound to an Agent. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| node_execution_id | string | Optional workflow node execution ID. When omitted, the latest active session for the node is used. | No | +| node_execution_id | string | Workflow node execution ID | Yes | | path | string | File path relative to the sandbox workspace | Yes | #### WorkflowAppLogPaginationResponse diff --git a/api/services/agent/composer_service.py b/api/services/agent/composer_service.py index aba9e9d8fc8..7e1b7697379 100644 --- a/api/services/agent/composer_service.py +++ b/api/services/agent/composer_service.py @@ -9,7 +9,7 @@ from sqlalchemy.sql.elements import ColumnElement from core.agent.publish_visibility import agent_has_workflow_callable_active_snapshot from libs.helper import to_timestamp -from models import Account +from models import Account, Conversation from models.agent import ( APP_BACKED_AGENT_SOURCES, Agent, @@ -18,12 +18,15 @@ from models.agent import ( AgentConfigRevision, AgentConfigRevisionOperation, AgentConfigSnapshot, + AgentConfigVersionKind, + AgentDebugConversation, AgentDriveFile, AgentIconType, AgentKind, AgentScope, AgentSource, AgentStatus, + AgentWorkspaceOwnerType, WorkflowAgentBindingType, WorkflowAgentNodeBinding, ) @@ -35,6 +38,7 @@ from models.workflow import Workflow from services.agent.agent_soul_state import agent_soul_has_model from services.agent.composer_validator import ComposerConfigValidator from services.agent.errors import ( + AgentBuildSandboxNotFoundError, AgentModelNotConfiguredError, AgentNameConflictError, AgentNotFoundError, @@ -42,11 +46,17 @@ from services.agent.errors import ( AgentVersionNotFoundError, InvalidComposerConfigError, ) +from services.agent.home_snapshot_service import ( + AgentHomeSnapshotService, + validate_home_snapshot_binding, +) from services.agent.knowledge_datasets import ( get_tenant_knowledge_dataset_rows, list_missing_tenant_knowledge_dataset_ids, ) +from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.roster_service import AgentRosterService +from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope from services.app_service import AppService, CreateAppParams from services.entities.agent_entities import ( AgentSoulConfig, @@ -56,6 +66,7 @@ from services.entities.agent_entities import ( ComposerVariant, WorkflowNodeJobConfig, ) +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection # WorkflowAgentNodeBinding.workflow_version tag for the draft workflow row. # Mirrors Workflow.version when it is "draft" (see models/workflow.py). @@ -198,6 +209,18 @@ class AgentComposerService: binding = cls._get_workflow_binding( session=session, tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id ) + retirement_candidates = ( + {binding.agent_id} + if binding is not None + and binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT + and binding.agent_id + and payload.save_strategy + in { + ComposerSaveStrategy.SAVE_AS_NEW_AGENT, + ComposerSaveStrategy.SAVE_TO_ROSTER, + } + else set() + ) match payload.save_strategy: case ComposerSaveStrategy.NODE_JOB_ONLY: @@ -232,7 +255,11 @@ class AgentComposerService: ) case ComposerSaveStrategy.SAVE_TO_ROSTER: binding = cls._save_to_roster( - session=session, tenant_id=tenant_id, account_id=account_id, binding=binding, payload=payload + session=session, + tenant_id=tenant_id, + account_id=account_id, + binding=binding, + payload=payload, ) session.flush() @@ -257,6 +284,17 @@ class AgentComposerService: payload=payload, agent_id=binding.agent_id, ) + session.commit() + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=tenant_id, + agent_ids=retirement_candidates, + account_id=account_id, + ) + enqueue_agent_resource_collection( + tenant_id=tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) return state @classmethod @@ -390,12 +428,10 @@ class AgentComposerService: @classmethod def _load_agent_composer_for_agent(cls, *, session: Session, tenant_id: str, agent: Agent) -> dict[str, Any]: - draft = cls._get_or_create_agent_draft( + draft = cls.get_or_create_normal_agent_draft( session=session, tenant_id=tenant_id, agent=agent, - draft_type=AgentConfigDraftType.DRAFT, - account_id=None, created_by=agent.updated_by or agent.created_by, ) version = cls._get_version_if_present( @@ -417,7 +453,34 @@ class AgentComposerService: @classmethod def save_agent_app_composer( - cls, *, session: Session, tenant_id: str, app_id: str, account_id: str, payload: ComposerSavePayload + cls, + *, + session: Session, + tenant_id: str, + app_id: str, + account_id: str, + payload: ComposerSavePayload, + ) -> dict[str, Any]: + try: + return cls._save_agent_app_composer_impl( + session=session, + tenant_id=tenant_id, + app_id=app_id, + account_id=account_id, + payload=payload, + ) + except IntegrityError as exc: + raise AgentNameConflictError() from exc + + @classmethod + def _save_agent_app_composer_impl( + cls, + *, + session: Session, + tenant_id: str, + app_id: str, + account_id: str, + payload: ComposerSavePayload, ) -> dict[str, Any]: if payload.variant != ComposerVariant.AGENT_APP: raise ValueError("Agent App composer endpoint only accepts agent_app variant") @@ -446,11 +509,25 @@ class AgentComposerService: updated_by=account_id, ) session.add(agent) - try: - session.flush() - except IntegrityError as exc: - session.rollback() - raise AgentNameConflictError() from exc + session.flush() + home_snapshot = AgentHomeSnapshotService.create_initial( + session=session, + tenant_id=tenant_id, + agent_id=agent.id, + ) + initial_version = cls._create_config_version( + session=session, + tenant_id=tenant_id, + agent_id=agent.id, + account_id=account_id, + agent_soul=AgentSoulConfig(), + operation=AgentConfigRevisionOperation.CREATE_VERSION, + version_note=None, + home_snapshot_id=home_snapshot.id, + ) + agent.active_config_snapshot_id = initial_version.id + agent.active_config_has_model = False + agent.active_config_is_published = False return cls._save_agent_composer_for_agent( session=session, tenant_id=tenant_id, @@ -488,7 +565,7 @@ class AgentComposerService: ) -> dict[str, Any]: if payload.agent_soul is None: raise ValueError("agent_soul is required") - cls._save_agent_draft( + draft = cls._save_agent_draft( session=session, tenant_id=tenant_id, agent=agent, @@ -503,6 +580,7 @@ class AgentComposerService: tenant_id=tenant_id, agent=agent, agent_soul=payload.agent_soul, + home_snapshot_id=draft.home_snapshot_id, ) session.flush() @@ -523,6 +601,7 @@ class AgentComposerService: tenant_id: str, agent: Agent, agent_soul: AgentSoulConfig, + home_snapshot_id: str, ) -> bool: if not agent.active_config_snapshot_id: return False @@ -538,7 +617,9 @@ class AgentComposerService: if not agent_has_workflow_callable_active_snapshot(session=session, agent=agent): return False - return _agent_soul_config_json(agent_soul) == _agent_soul_config_json(active_version.config_snapshot_dict) + return home_snapshot_id == active_version.home_snapshot_id and _agent_soul_config_json( + agent_soul + ) == _agent_soul_config_json(active_version.config_snapshot_dict) @classmethod def publish_agent_app_draft( @@ -567,6 +648,11 @@ class AgentComposerService: if not agent_soul_has_model(agent_soul): raise AgentModelNotConfiguredError() cls.validate_knowledge_datasets(session=session, tenant_id=tenant_id, agent_soul=agent_soul) + validate_home_snapshot_binding( + session=session, + agent=agent, + home_snapshot_id=draft.home_snapshot_id, + ) version = cls._create_config_version( session=session, tenant_id=tenant_id, @@ -576,6 +662,7 @@ class AgentComposerService: operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, version_note=version_note, previous_snapshot_id=agent.active_config_snapshot_id, + home_snapshot_id=draft.home_snapshot_id, ) agent.active_config_snapshot_id = version.id agent.active_config_has_model = agent_soul_has_model(agent_soul) @@ -595,6 +682,29 @@ class AgentComposerService: def checkout_agent_app_build_draft( cls, *, session: Session, tenant_id: str, agent_id: str, account_id: str, force: bool = False ) -> dict[str, Any]: + try: + result, retired_binding_id = cls._checkout_agent_app_build_draft_in_transaction( + session=session, + tenant_id=tenant_id, + agent_id=agent_id, + account_id=account_id, + force=force, + ) + session.commit() + except Exception: + session.rollback() + raise + if retired_binding_id is not None: + enqueue_agent_resource_collection( + tenant_id=tenant_id, + binding_ids=(retired_binding_id,), + ) + return result + + @classmethod + def _checkout_agent_app_build_draft_in_transaction( + cls, *, session: Session, tenant_id: str, agent_id: str, account_id: str, force: bool + ) -> tuple[dict[str, Any], str | None]: agent = cls._require_agent(session=session, tenant_id=tenant_id, agent_id=agent_id) normal_draft = cls._get_or_create_agent_draft( session=session, @@ -612,7 +722,23 @@ class AgentComposerService: account_id=account_id, ) if build_draft is not None and not force: - return cls._serialize_build_draft_state(build_draft) + return cls._serialize_build_draft_state(build_draft), None + retired_binding_id: str | None = None + if build_draft is not None and build_draft.agent_workspace_binding_id is not None: + cls._validate_active_build_draft_binding( + session=session, + tenant_id=tenant_id, + agent=agent, + build_draft=build_draft, + ) + retired_binding_id = AgentWorkspaceService.retire_binding( + session=session, + tenant_id=tenant_id, + binding_id=build_draft.agent_workspace_binding_id, + ) + if retired_binding_id is None: + raise AgentBuildSandboxNotFoundError() + build_draft.agent_workspace_binding_id = None if build_draft is None: build_draft = AgentConfigDraft( tenant_id=tenant_id, @@ -624,10 +750,44 @@ class AgentComposerService: ) session.add(build_draft) build_draft.base_snapshot_id = normal_draft.base_snapshot_id + build_draft.home_snapshot_id = normal_draft.home_snapshot_id build_draft.config_snapshot = AgentSoulConfig.model_validate(normal_draft.config_snapshot_dict) build_draft.updated_by = account_id session.flush() - return cls._serialize_build_draft_state(build_draft) + return cls._serialize_build_draft_state(build_draft), retired_binding_id + + @classmethod + def _validate_active_build_draft_binding( + cls, + *, + session: Session, + tenant_id: str, + agent: Agent, + build_draft: AgentConfigDraft, + ) -> None: + binding_id = build_draft.agent_workspace_binding_id + runtime_app_id = AgentRosterService.runtime_backing_app_id(agent) + if binding_id is None or runtime_app_id is None: + raise AgentBuildSandboxNotFoundError() + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=tenant_id, + binding_id=binding_id, + expected_owner_scope=WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=runtime_app_id, + owner_type=AgentWorkspaceOwnerType.BUILD_DRAFT, + owner_id=build_draft.id, + ), + ) + if binding is None or binding.agent_id != agent.id: + raise AgentBuildSandboxNotFoundError() + AgentWorkspaceService.validate_binding_generation( + binding, + base_home_snapshot_id=build_draft.home_snapshot_id, + agent_config_version_id=build_draft.id, + agent_config_version_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) @classmethod def load_agent_app_build_draft( @@ -669,6 +829,29 @@ class AgentComposerService: def apply_agent_app_build_draft( cls, *, session: Session, tenant_id: str, agent_id: str, account_id: str ) -> dict[str, Any]: + try: + result, retired_binding_ids = cls._apply_agent_app_build_draft_in_transaction( + session=session, + tenant_id=tenant_id, + agent_id=agent_id, + account_id=account_id, + ) + session.commit() + except Exception: + session.rollback() + raise + enqueue_agent_resource_collection(tenant_id=tenant_id, binding_ids=retired_binding_ids) + return result + + @classmethod + def _apply_agent_app_build_draft_in_transaction( + cls, + *, + session: Session, + tenant_id: str, + agent_id: str, + account_id: str, + ) -> tuple[dict[str, Any], list[str]]: agent = cls._require_agent(session=session, tenant_id=tenant_id, agent_id=agent_id) build_draft = cls._get_agent_draft( session=session, @@ -680,6 +863,21 @@ class AgentComposerService: if build_draft is None: raise AgentVersionNotFoundError() applied_agent_soul = AgentSoulConfig.model_validate(build_draft.config_snapshot_dict) + ComposerConfigValidator.validate_publish_payload( + ComposerSavePayload( + variant=ComposerVariant.AGENT_APP, + agent_soul=applied_agent_soul, + save_strategy=ComposerSaveStrategy.SAVE_AS_NEW_VERSION, + ) + ) + cls.validate_knowledge_datasets(session=session, tenant_id=tenant_id, agent_soul=applied_agent_soul) + source_binding_id = build_draft.agent_workspace_binding_id + if source_binding_id is None: + raise AgentBuildSandboxNotFoundError() + home_snapshot = AgentHomeSnapshotService.create_for_build_apply( + session=session, + build_draft=build_draft, + ) normal_draft = cls._save_agent_draft( session=session, tenant_id=tenant_id, @@ -690,32 +888,145 @@ class AgentComposerService: account_id_for_audit=account_id, base_snapshot_id=build_draft.base_snapshot_id, ) + retired_binding_ids = cls._retire_normal_preview_bindings( + session=session, + tenant_id=tenant_id, + agent=agent, + normal_draft=normal_draft, + ) + normal_draft.home_snapshot_id = home_snapshot.id agent.active_config_is_published = cls._agent_soul_matches_active_config( session=session, tenant_id=tenant_id, agent=agent, agent_soul=applied_agent_soul, + home_snapshot_id=home_snapshot.id, ) agent.updated_by = account_id + retired_binding_id = AgentWorkspaceService.retire_binding( + session=session, + tenant_id=tenant_id, + binding_id=source_binding_id, + ) + if retired_binding_id is None: + raise AgentBuildSandboxNotFoundError() + retired_binding_ids.append(source_binding_id) session.delete(build_draft) - session.flush() - return {"result": "success", "draft": cls._serialize_draft(normal_draft)} + return {"result": "success", "draft": cls._serialize_draft(normal_draft)}, retired_binding_ids + + @classmethod + def _retire_normal_preview_bindings( + cls, + *, + session: Session, + tenant_id: str, + agent: Agent, + normal_draft: AgentConfigDraft, + ) -> list[str]: + """Retire Preview participants before Build Apply replaces the shared Draft Home.""" + + mappings = session.scalars( + select(AgentDebugConversation).where( + AgentDebugConversation.tenant_id == tenant_id, + AgentDebugConversation.agent_id == agent.id, + AgentDebugConversation.draft_type == AgentConfigDraftType.DRAFT, + ) + ).all() + retired_binding_ids: list[str] = [] + for mapping in mappings: + conversation = session.scalar( + select(Conversation).where( + Conversation.id == mapping.conversation_id, + Conversation.app_id == mapping.app_id, + Conversation.is_deleted.is_(False), + ) + ) + if conversation is None or conversation.agent_workspace_binding_id is None: + continue + binding_id = conversation.agent_workspace_binding_id + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=tenant_id, + binding_id=binding_id, + expected_owner_scope=WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=mapping.app_id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=conversation.id, + ), + ) + if binding is None or binding.agent_id != agent.id: + raise AgentWorkspaceNotFoundError("Agent Preview participant Binding is unavailable") + AgentWorkspaceService.validate_binding_generation( + binding, + base_home_snapshot_id=normal_draft.home_snapshot_id, + agent_config_version_id=normal_draft.id, + agent_config_version_kind=AgentConfigVersionKind.DRAFT, + ) + retired_binding_id = AgentWorkspaceService.retire_binding( + session=session, + tenant_id=tenant_id, + binding_id=binding_id, + ) + if retired_binding_id is None: + raise AgentWorkspaceNotFoundError("Agent Preview participant Binding is unavailable") + conversation.agent_workspace_binding_id = None + retired_binding_ids.append(binding_id) + return retired_binding_ids @classmethod def discard_agent_app_build_draft( cls, *, session: Session, tenant_id: str, agent_id: str, account_id: str ) -> dict[str, Any]: + try: + result, retired_binding_id = cls._discard_agent_app_build_draft_in_transaction( + session=session, + tenant_id=tenant_id, + agent_id=agent_id, + account_id=account_id, + ) + session.commit() + except Exception: + session.rollback() + raise + if retired_binding_id is not None: + enqueue_agent_resource_collection( + tenant_id=tenant_id, + binding_ids=(retired_binding_id,), + ) + return result + + @classmethod + def _discard_agent_app_build_draft_in_transaction( + cls, *, session: Session, tenant_id: str, agent_id: str, account_id: str + ) -> tuple[dict[str, Any], str | None]: + agent = cls._require_agent(session=session, tenant_id=tenant_id, agent_id=agent_id) build_draft = cls._get_agent_draft( session=session, tenant_id=tenant_id, - agent_id=agent_id, + agent_id=agent.id, draft_type=AgentConfigDraftType.DEBUG_BUILD, account_id=account_id, ) - if build_draft is not None: - session.delete(build_draft) - session.flush() - return {"result": "success"} + if build_draft is None: + return {"result": "success"}, None + retired_binding_id: str | None = None + if build_draft.agent_workspace_binding_id is not None: + cls._validate_active_build_draft_binding( + session=session, + tenant_id=tenant_id, + agent=agent, + build_draft=build_draft, + ) + retired_binding_id = AgentWorkspaceService.retire_binding( + session=session, + tenant_id=tenant_id, + binding_id=build_draft.agent_workspace_binding_id, + ) + if retired_binding_id is None: + raise AgentBuildSandboxNotFoundError() + session.delete(build_draft) + return {"result": "success"}, retired_binding_id @classmethod def collect_validation_findings( @@ -1279,6 +1590,12 @@ class AgentComposerService: binding = cls._require_binding(binding) if not binding.agent_id or payload.agent_soul is None: raise ValueError("agent_id and agent_soul are required") + current_snapshot = cls._require_version( + session=session, + tenant_id=tenant_id, + agent_id=binding.agent_id, + version_id=binding.current_snapshot_id, + ) version = cls._create_config_version( session=session, tenant_id=tenant_id, @@ -1287,6 +1604,7 @@ class AgentComposerService: agent_soul=payload.agent_soul, operation=AgentConfigRevisionOperation.SAVE_NEW_VERSION, version_note=payload.version_note, + home_snapshot_id=current_snapshot.home_snapshot_id, ) agent = cls._require_agent(session=session, tenant_id=tenant_id, agent_id=binding.agent_id) agent.active_config_snapshot_id = version.id @@ -1449,6 +1767,11 @@ class AgentComposerService: ) session.add(agent) session.flush() + home_snapshot = AgentHomeSnapshotService.create_initial( + session=session, + tenant_id=tenant_id, + agent_id=agent.id, + ) version = cls._create_config_version( session=session, tenant_id=tenant_id, @@ -1457,6 +1780,7 @@ class AgentComposerService: agent_soul=agent_soul, operation=AgentConfigRevisionOperation.CREATE_VERSION, version_note=None, + home_snapshot_id=home_snapshot.id, ) agent.active_config_snapshot_id = version.id agent.active_config_has_model = agent_soul_has_model(agent_soul) @@ -1590,7 +1914,6 @@ class AgentComposerService: session=session, ) except IntegrityError as exc: - session.rollback() raise AgentNameConflictError() from exc agent = AgentRosterService(session).get_app_backing_agent(tenant_id=tenant_id, app_id=app.id) @@ -1628,6 +1951,7 @@ class AgentComposerService: agent_soul: AgentSoulConfig, operation: AgentConfigRevisionOperation, version_note: str | None, + home_snapshot_id: str, previous_snapshot_id: str | None = None, ) -> AgentConfigSnapshot: next_version = ( @@ -1644,6 +1968,7 @@ class AgentComposerService: agent_id=agent_id, version=next_version, config_snapshot=agent_soul, + home_snapshot_id=home_snapshot_id, version_note=version_note, created_by=account_id, ) @@ -1683,6 +2008,7 @@ class AgentComposerService: operation=operation, version_note=version_note, previous_snapshot_id=current_snapshot.id, + home_snapshot_id=current_snapshot.home_snapshot_id, ) @classmethod @@ -1749,7 +2075,10 @@ class AgentComposerService: agent: Agent, created_by: str | None, ) -> AgentConfigDraft: - """Resolve the shared Preview draft, rebasing inline agents when needed.""" + """Resolve the normal Draft, rebasing only stale WORKFLOW_ONLY DRAFT rows whose account_id is None. + + Roster and DEBUG_BUILD Drafts are never rebased. + """ return cls._get_or_create_agent_draft( session=session, tenant_id=tenant_id, @@ -1767,6 +2096,8 @@ class AgentComposerService: snapshot: AgentConfigSnapshot, updated_by: str | None, ) -> bool: + """Sync a stale normal Draft's base_snapshot_id, home_snapshot_id, config_snapshot, and updated_by.""" + if ( agent.scope != AgentScope.WORKFLOW_ONLY or draft.draft_type != AgentConfigDraftType.DRAFT @@ -1777,6 +2108,7 @@ class AgentComposerService: ): return False draft.base_snapshot_id = snapshot.id + draft.home_snapshot_id = snapshot.home_snapshot_id draft.config_snapshot = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) draft.updated_by = updated_by return True @@ -1813,7 +2145,9 @@ class AgentComposerService: agent_id=agent.id, version_id=agent.active_config_snapshot_id, ) - if active_snapshot is not None and cls._rebase_workflow_only_normal_draft( + if active_snapshot is None: + raise AgentVersionNotFoundError() + if cls._rebase_workflow_only_normal_draft( agent=agent, draft=draft, snapshot=active_snapshot, @@ -1827,18 +2161,17 @@ class AgentComposerService: agent_id=agent.id, version_id=agent.active_config_snapshot_id, ) - agent_soul = ( - AgentSoulConfig.model_validate(base_snapshot.config_snapshot_dict) - if base_snapshot is not None - else AgentSoulConfig() - ) + if base_snapshot is None: + raise AgentVersionNotFoundError() + agent_soul = AgentSoulConfig.model_validate(base_snapshot.config_snapshot_dict) draft = AgentConfigDraft( tenant_id=tenant_id, agent_id=agent.id, draft_type=draft_type, account_id=account_id if draft_type == AgentConfigDraftType.DEBUG_BUILD else None, draft_owner_key=account_id if draft_type == AgentConfigDraftType.DEBUG_BUILD and account_id else "", - base_snapshot_id=base_snapshot.id if base_snapshot else None, + base_snapshot_id=base_snapshot.id, + home_snapshot_id=base_snapshot.home_snapshot_id, config_snapshot=agent_soul, created_by=created_by, updated_by=created_by, @@ -2140,7 +2473,7 @@ class AgentComposerService: from services.agent.roster_service import AgentRosterService - return AgentRosterService(session).get_or_create_agent_app_debug_conversation_id( + return AgentRosterService(session).get_or_create_build_conversation( tenant_id=tenant_id, agent_id=agent.id, account_id=account_id, diff --git a/api/services/agent/dsl_service.py b/api/services/agent/dsl_service.py index b2ffae45dad..b45ff749a52 100644 --- a/api/services/agent/dsl_service.py +++ b/api/services/agent/dsl_service.py @@ -49,6 +49,7 @@ from services.agent.dsl_entities import ( make_portable_agent_package, portable_ref, ) +from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.knowledge_datasets import get_tenant_knowledge_dataset_rows from services.agent.roster_service import AgentRosterService from services.entities.dsl_entities import DslImportWarning @@ -224,6 +225,7 @@ class AgentDslService: account_id=None, draft_owner_key="", base_snapshot_id=snapshot.id, + home_snapshot_id=snapshot.home_snapshot_id, config_snapshot=soul, created_by=account.id, updated_by=account.id, @@ -243,7 +245,7 @@ class AgentDslService: portable_graph: Mapping[str, Any], raw_packages: Mapping[str, Any], account: Account, - ) -> tuple[dict[str, Any], list[DslImportWarning]]: + ) -> tuple[dict[str, Any], list[DslImportWarning], set[str]]: """Materialize every packaged Agent as a node-owned inline Agent.""" graph = copy.deepcopy(dict(portable_graph)) @@ -256,6 +258,11 @@ class AgentDslService: WorkflowAgentNodeBinding.workflow_version == Workflow.VERSION_DRAFT, ) ).all() + retirement_candidates = { + binding.agent_id + for binding in previous_bindings + if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id + } for binding in previous_bindings: self.session.delete(binding) self.session.flush() @@ -312,7 +319,7 @@ class AgentDslService: workflow.graph = json.dumps(graph) self.session.flush() - return graph, warnings + return graph, warnings, retirement_candidates def clone_inline_binding_for_node( self, @@ -562,11 +569,17 @@ class AgentDslService: ) or 0 ) + 1 + home_snapshot = AgentHomeSnapshotService.create_initial( + session=self.session, + tenant_id=tenant_id, + agent_id=agent.id, + ) snapshot = AgentConfigSnapshot( tenant_id=tenant_id, agent_id=agent.id, version=next_version, config_snapshot=soul, + home_snapshot_id=home_snapshot.id, created_by=account_id, ) self.session.add(snapshot) diff --git a/api/services/agent/errors.py b/api/services/agent/errors.py index 163687815d8..b900bd857bb 100644 --- a/api/services/agent/errors.py +++ b/api/services/agent/errors.py @@ -29,6 +29,12 @@ class AgentModelNotConfiguredError(BaseHTTPException): code = 400 +class AgentBuildSandboxNotFoundError(BaseHTTPException): + error_code = "agent_build_sandbox_not_found" + description = "The retained Build Sandbox is no longer available." + code = 404 + + class AgentSoulLockedError(BadRequest): description = "Agent Soul is locked for this workflow node." diff --git a/api/services/agent/home_snapshot_service.py b/api/services/agent/home_snapshot_service.py new file mode 100644 index 00000000000..6f4fb444480 --- /dev/null +++ b/api/services/agent/home_snapshot_service.py @@ -0,0 +1,238 @@ +"""Own immutable Agent Home Snapshot ledger rows and physical collection.""" + +from __future__ import annotations + +import logging + +from dify_agent.client import Client, DifyAgentNotFoundError +from dify_agent.protocol import CreateHomeSnapshotFromBindingRequest, InitializeHomeSnapshotRequest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from configs import dify_config +from core.db.session_factory import session_factory +from libs.datetime_utils import naive_utc_now +from libs.uuid_utils import uuidv7 +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigSnapshot, + AgentConfigVersionKind, + AgentHomeSnapshot, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceOwnerType, +) +from services.agent.errors import AgentBuildSandboxNotFoundError +from services.agent.workspace_service import AgentWorkspaceService, WorkspaceOwnerScope + +logger = logging.getLogger(__name__) + + +class AgentHomeSnapshotUnavailableError(RuntimeError): + """The requested owner-scoped Home Snapshot cannot be used.""" + + +class AgentHomeSnapshotService: + """Create, retire, and collect Agent-owned immutable Home Snapshots.""" + + @classmethod + def create_initial( + cls, + *, + session: Session, + tenant_id: str, + agent_id: str, + ) -> AgentHomeSnapshot: + home_snapshot_id = str(uuidv7()) + with cls._client() as client: + response = client.initialize_home_snapshot_sync( + InitializeHomeSnapshotRequest( + tenant_id=tenant_id, + agent_id=agent_id, + home_snapshot_id=home_snapshot_id, + ) + ) + home_snapshot = AgentHomeSnapshot( + id=home_snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + snapshot_ref=response.snapshot_ref, + status=AgentWorkingResourceStatus.ACTIVE, + ) + session.add(home_snapshot) + session.flush() + return home_snapshot + + @classmethod + def create_for_build_apply( + cls, + *, + session: Session, + build_draft: AgentConfigDraft, + ) -> AgentHomeSnapshot: + """Checkpoint the exact participant owned by ``build_draft``.""" + + source_binding_id = build_draft.agent_workspace_binding_id + if source_binding_id is None: + raise AgentBuildSandboxNotFoundError() + agent = session.scalar( + select(Agent).where( + Agent.id == build_draft.agent_id, + Agent.tenant_id == build_draft.tenant_id, + ) + ) + if agent is None: + raise AgentBuildSandboxNotFoundError() + from services.agent.roster_service import AgentRosterService + + runtime_app_id = AgentRosterService.runtime_backing_app_id(agent) + if runtime_app_id is None: + raise AgentBuildSandboxNotFoundError() + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=build_draft.tenant_id, + binding_id=source_binding_id, + expected_owner_scope=WorkspaceOwnerScope( + tenant_id=build_draft.tenant_id, + app_id=runtime_app_id, + owner_type=AgentWorkspaceOwnerType.BUILD_DRAFT, + owner_id=build_draft.id, + ), + ) + if binding is None or binding.agent_id != build_draft.agent_id: + raise AgentBuildSandboxNotFoundError() + AgentWorkspaceService.validate_binding_generation( + binding, + base_home_snapshot_id=build_draft.home_snapshot_id, + agent_config_version_id=build_draft.id, + agent_config_version_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) + + home_snapshot_id = str(uuidv7()) + try: + with cls._client() as client: + response = client.create_home_snapshot_from_binding_sync( + CreateHomeSnapshotFromBindingRequest( + tenant_id=build_draft.tenant_id, + agent_id=build_draft.agent_id, + home_snapshot_id=home_snapshot_id, + backend_binding_ref=binding.backend_binding_ref, + ) + ) + except DifyAgentNotFoundError as exc: + raise AgentBuildSandboxNotFoundError() from exc + + home_snapshot = AgentHomeSnapshot( + id=home_snapshot_id, + tenant_id=build_draft.tenant_id, + agent_id=build_draft.agent_id, + snapshot_ref=response.snapshot_ref, + status=AgentWorkingResourceStatus.ACTIVE, + ) + session.add(home_snapshot) + return home_snapshot + + @classmethod + def retire_all_for_agent(cls, *, session: Session, tenant_id: str, agent_id: str) -> list[str]: + rows = session.scalars( + select(AgentHomeSnapshot).where( + AgentHomeSnapshot.tenant_id == tenant_id, + AgentHomeSnapshot.agent_id == agent_id, + AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE, + ) + ).all() + now = naive_utc_now() + for row in rows: + row.status = AgentWorkingResourceStatus.RETIRED + row.retired_at = now + return [row.id for row in rows] + + @classmethod + def collect_retired_home_snapshot(cls, *, tenant_id: str, home_snapshot_id: str) -> None: + try: + cls._collect_retired_home_snapshot(tenant_id=tenant_id, home_snapshot_id=home_snapshot_id) + except Exception: + logger.exception( + "Failed to collect retired Agent Home Snapshot", + extra={"tenant_id": tenant_id, "home_snapshot_id": home_snapshot_id}, + ) + + @classmethod + def _collect_retired_home_snapshot(cls, *, tenant_id: str, home_snapshot_id: str) -> None: + with session_factory.create_session() as session: + snapshot = session.scalar( + select(AgentHomeSnapshot).where( + AgentHomeSnapshot.id == home_snapshot_id, + AgentHomeSnapshot.tenant_id == tenant_id, + AgentHomeSnapshot.status == AgentWorkingResourceStatus.RETIRED, + ) + ) + if snapshot is None: + return + referenced = session.scalar( + select(AgentConfigDraft.id).where(AgentConfigDraft.home_snapshot_id == home_snapshot_id).limit(1) + ) or session.scalar( + select(AgentConfigSnapshot.id).where(AgentConfigSnapshot.home_snapshot_id == home_snapshot_id).limit(1) + ) + if referenced is not None: + return + snapshot_ref = snapshot.snapshot_ref + try: + cls.delete(snapshot_ref=snapshot_ref) + except Exception: + logger.exception( + "Failed to collect retired Agent Home Snapshot", + extra={"tenant_id": tenant_id, "home_snapshot_id": home_snapshot_id}, + ) + return + with session_factory.create_session() as session: + snapshot = session.scalar( + select(AgentHomeSnapshot).where( + AgentHomeSnapshot.id == home_snapshot_id, + AgentHomeSnapshot.tenant_id == tenant_id, + AgentHomeSnapshot.status == AgentWorkingResourceStatus.RETIRED, + ) + ) + if snapshot is not None: + session.delete(snapshot) + session.commit() + + @classmethod + def delete(cls, *, snapshot_ref: str) -> None: + with cls._client() as client: + client.delete_home_snapshot_sync(snapshot_ref) + + @staticmethod + def _client() -> Client: + base_url = dify_config.AGENT_BACKEND_BASE_URL + if not base_url: + raise AgentHomeSnapshotUnavailableError("Dify Agent backend is required for Home Snapshot operations") + return Client(base_url=base_url) + + +def validate_home_snapshot_binding(*, session: Session, agent: Agent, home_snapshot_id: str) -> None: + _require_owned_home_snapshot(session=session, agent=agent, home_snapshot_id=home_snapshot_id) + + +def _require_owned_home_snapshot(*, session: Session, agent: Agent, home_snapshot_id: str) -> AgentHomeSnapshot: + if agent.status != AgentStatus.ACTIVE: + raise AgentHomeSnapshotUnavailableError(f"Agent {agent.id} is not active") + home_snapshot = session.scalar( + select(AgentHomeSnapshot).where( + AgentHomeSnapshot.id == home_snapshot_id, + AgentHomeSnapshot.tenant_id == agent.tenant_id, + AgentHomeSnapshot.agent_id == agent.id, + AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE, + ) + ) + if home_snapshot is None: + raise AgentHomeSnapshotUnavailableError(f"Home Snapshot {home_snapshot_id} is unavailable for Agent {agent.id}") + return home_snapshot + + +__all__ = [ + "AgentHomeSnapshotService", + "AgentHomeSnapshotUnavailableError", + "validate_home_snapshot_binding", +] diff --git a/api/services/agent/retirement_service.py b/api/services/agent/retirement_service.py new file mode 100644 index 00000000000..25b133e14f9 --- /dev/null +++ b/api/services/agent/retirement_service.py @@ -0,0 +1,166 @@ +"""Workflow-only Agent ownership retirement after product transactions commit.""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable + +from sqlalchemy import or_, select +from sqlalchemy.orm import Session + +from core.db.session_factory import session_factory +from libs.datetime_utils import naive_utc_now +from models.agent import ( + Agent, + AgentScope, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, + WorkflowAgentNodeBinding, +) +from models.enums import AppStatus +from models.model import App +from models.workflow import Workflow +from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.workspace_service import AgentWorkspaceService + +logger = logging.getLogger(__name__) + + +class WorkflowAgentRetirementService: + """Archive workflow-only Agents once no effective binding owns them.""" + + @classmethod + def retire_unowned( + cls, + *, + tenant_id: str, + agent_ids: Iterable[str], + account_id: str | None, + ) -> tuple[list[str], list[str]]: + """Re-check ownership, archive orphans, and commit their resource retirement.""" + + candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) + if not candidates: + return [], [] + retired_bindings: list[str] = [] + retired_snapshots: list[str] = [] + try: + with session_factory.create_session() as session: + retired_agent_ids = cls.archive_unowned( + session=session, + tenant_id=tenant_id, + agent_ids=candidates, + account_id=account_id, + ) + for agent_id in retired_agent_ids: + bindings = session.scalars( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.agent_id == agent_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + ).all() + for binding in bindings: + binding_id = AgentWorkspaceService.retire_binding( + session=session, + tenant_id=tenant_id, + binding_id=binding.id, + ) + if binding_id is not None: + retired_bindings.append(binding_id) + retired_snapshots.extend( + AgentHomeSnapshotService.retire_all_for_agent( + session=session, + tenant_id=tenant_id, + agent_id=agent_id, + ) + ) + session.commit() + except Exception: + logger.exception( + "Failed to retire unowned Workflow Agents", + extra={ + "tenant_id": tenant_id, + "agent_ids": candidates, + }, + ) + return [], [] + return retired_bindings, retired_snapshots + + @classmethod + def archive_unowned( + cls, + *, + session: Session, + tenant_id: str, + agent_ids: Iterable[str], + account_id: str | None, + ) -> list[str]: + """Archive active orphans and return every orphan eligible for Home cleanup.""" + candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) + if not candidates: + return [] + agents = session.scalars( + select(Agent).where( + Agent.tenant_id == tenant_id, + Agent.id.in_(candidates), + Agent.scope == AgentScope.WORKFLOW_ONLY, + Agent.status.in_((AgentStatus.ACTIVE, AgentStatus.ARCHIVED)), + ) + ).all() + effective_agent_ids = cls._effective_agent_ids( + session=session, + tenant_id=tenant_id, + agent_ids=[agent.id for agent in agents], + ) + now = naive_utc_now() + cleanup_candidates: list[str] = [] + for agent in agents: + if agent.id in effective_agent_ids: + continue + if agent.status == AgentStatus.ACTIVE: + agent.status = AgentStatus.ARCHIVED + agent.archived_by = account_id + agent.archived_at = now + agent.updated_by = account_id or agent.updated_by + agent.updated_at = now + cleanup_candidates.append(agent.id) + session.flush() + return cleanup_candidates + + @staticmethod + def _effective_agent_ids( + *, + session: Session, + tenant_id: str, + agent_ids: list[str], + ) -> set[str]: + if not agent_ids: + return set() + values = session.scalars( + select(WorkflowAgentNodeBinding.agent_id) + .join( + Workflow, + Workflow.id == WorkflowAgentNodeBinding.workflow_id, + ) + .join(App, App.id == WorkflowAgentNodeBinding.app_id) + .where( + WorkflowAgentNodeBinding.tenant_id == tenant_id, + WorkflowAgentNodeBinding.agent_id.in_(agent_ids), + Workflow.tenant_id == tenant_id, + Workflow.app_id == WorkflowAgentNodeBinding.app_id, + Workflow.version == WorkflowAgentNodeBinding.workflow_version, + App.tenant_id == tenant_id, + App.status == AppStatus.NORMAL, + or_( + Workflow.version == Workflow.VERSION_DRAFT, + App.workflow_id == Workflow.id, + ), + ) + .distinct() + ).all() + return {agent_id for agent_id in values if agent_id} + + +__all__ = ["WorkflowAgentRetirementService"] diff --git a/api/services/agent/roster_service.py b/api/services/agent/roster_service.py index c366f2af47b..37eb03cf32b 100644 --- a/api/services/agent/roster_service.py +++ b/api/services/agent/roster_service.py @@ -4,10 +4,8 @@ from typing import Any, TypedDict from sqlalchemy import and_, func, or_, select from sqlalchemy.exc import IntegrityError -from clients.agent_backend.session_cleanup import AgentBackendSessionCleanupPayload from constants.model_template import default_app_templates from core.agent.publish_visibility import workflow_callable_active_snapshot_filter -from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore from core.app.entities.app_invoke_entities import InvokeFrom from libs.datetime_utils import naive_utc_now from libs.helper import to_timestamp @@ -25,6 +23,9 @@ from models.agent import ( AgentScope, AgentSource, AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, WorkflowAgentBindingType, WorkflowAgentNodeBinding, ) @@ -36,15 +37,18 @@ from services.agent.agent_soul_state import agent_soul_has_model from services.agent.composer_validator import ComposerConfigValidator from services.agent.errors import ( AgentArchivedError, + AgentBuildSandboxNotFoundError, AgentNameConflictError, AgentNotFoundError, AgentVersionNotFoundError, ) +from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope from services.app_service import AppService, CreateAppParams from services.enterprise.enterprise_service import EnterpriseService from services.entities.agent_entities import RosterAgentCreatePayload, RosterAgentUpdatePayload from services.feature_service import FeatureService -from tasks.agent_backend_session_cleanup_task import cleanup_conversation_agent_runtime_session +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -304,6 +308,27 @@ class AgentRosterService: account_id: str, payload: RosterAgentCreatePayload, source: AgentSource = AgentSource.ROSTER, + ) -> Agent: + try: + agent = self._create_roster_agent_in_transaction( + tenant_id=tenant_id, + account_id=account_id, + payload=payload, + source=source, + ) + self._session.commit() + return agent + except IntegrityError as exc: + self._session.rollback() + raise AgentNameConflictError() from exc + + def _create_roster_agent_in_transaction( + self, + *, + tenant_id: str, + account_id: str, + payload: RosterAgentCreatePayload, + source: AgentSource, ) -> Agent: ComposerConfigValidator.validate_agent_soul(payload.agent_soul) @@ -323,17 +348,19 @@ class AgentRosterService: updated_by=account_id, ) self._session.add(agent) - try: - self._session.flush() - except IntegrityError as exc: - self._session.rollback() - raise AgentNameConflictError() from exc + self._session.flush() + home_snapshot = AgentHomeSnapshotService.create_initial( + session=self._session, + tenant_id=tenant_id, + agent_id=agent.id, + ) version = AgentConfigSnapshot( tenant_id=tenant_id, agent_id=agent.id, version=1, config_snapshot=payload.agent_soul, + home_snapshot_id=home_snapshot.id, version_note=payload.version_note, created_by=account_id, ) @@ -354,11 +381,6 @@ class AgentRosterService: agent.active_config_has_model = agent_soul_has_model(payload.agent_soul) agent.active_config_is_published = True - try: - self._session.commit() - except IntegrityError as exc: - self._session.rollback() - raise AgentNameConflictError() from exc return agent def create_backing_agent_for_app( @@ -405,17 +427,19 @@ class AgentRosterService: updated_by=account_id, ) self._session.add(agent) - try: - self._session.flush() - except IntegrityError as exc: - self._session.rollback() - raise AgentNameConflictError() from exc + self._session.flush() + home_snapshot = AgentHomeSnapshotService.create_initial( + session=self._session, + tenant_id=tenant_id, + agent_id=agent.id, + ) version = AgentConfigSnapshot( tenant_id=tenant_id, agent_id=agent.id, version=1, config_snapshot=soul, + home_snapshot_id=home_snapshot.id, created_by=account_id, ) self._session.add(version) @@ -590,16 +614,15 @@ class AgentRosterService: self._session.flush() return conversation_id - def get_or_create_agent_app_debug_conversation_id( + def get_or_create_build_conversation( self, *, tenant_id: str, agent_id: str, account_id: str, - draft_type: AgentConfigDraftType = AgentConfigDraftType.DEBUG_BUILD, commit: bool = True, ) -> str: - """Return the current editor's Build or Preview conversation for an Agent App.""" + """Return the current editor's stable Build conversation.""" agent = self._session.scalar( select(Agent).where( @@ -614,21 +637,20 @@ class AgentRosterService: conversation_id = self._get_or_create_agent_app_debug_conversation( agent=agent, account_id=account_id, - draft_type=draft_type, + draft_type=AgentConfigDraftType.DEBUG_BUILD, ) if commit: self._session.commit() return conversation_id - def load_agent_app_debug_conversation_id( + def get_current_preview_conversation( self, *, tenant_id: str, agent_id: str, account_id: str, - draft_type: AgentConfigDraftType = AgentConfigDraftType.DEBUG_BUILD, ) -> str | None: - """Return the editor's existing scoped conversation without creating or repairing rows.""" + """Return the editor's current Preview conversation without creating one.""" return self._session.scalar( select(Conversation.id) @@ -637,7 +659,7 @@ class AgentRosterService: AgentDebugConversation.tenant_id == tenant_id, AgentDebugConversation.agent_id == agent_id, AgentDebugConversation.account_id == account_id, - AgentDebugConversation.draft_type == draft_type, + AgentDebugConversation.draft_type == AgentConfigDraftType.DRAFT, AgentDebugConversation.app_id == Conversation.app_id, Conversation.from_source == ConversationFromSource.CONSOLE, Conversation.from_account_id == account_id, @@ -657,25 +679,12 @@ class AgentRosterService: or 0 ) - def refresh_agent_app_debug_conversation_id( - self, - *, - tenant_id: str, - agent_id: str, - account_id: str, - draft_type: AgentConfigDraftType = AgentConfigDraftType.DEBUG_BUILD, - ) -> str: - """Start a new scoped console conversation for the current Agent App editor. + def rotate_preview_conversation(self, *, tenant_id: str, agent_id: str, account_id: str) -> str: + """Rotate Preview and retire its exact Conversation-owned Binding. - If this account already has a mapping for the requested draft surface, the previous - conversation is abandoned after the replacement mapping is committed: any ACTIVE - conversation-owned Agent runtime sessions for that old conversation are sent through - best-effort backend cleanup and then retired locally even when enqueueing fails. This - order prevents a failed database commit from retiring the still-current runtime session. - The other draft surface is left untouched. - - A user and draft surface own one current mapping. If new-conversation requests overlap, - the last committed rotation becomes current and earlier response IDs cannot be continued. + The mapping update and exact CONVERSATION Binding retirement commit in + one transaction. Validation failures fail fast; collection is enqueued + only after commit. """ agent = self._session.scalar( @@ -694,142 +703,194 @@ class AgentRosterService: if not backing_app_id: raise AgentNotFoundError() - conversation_id = self._create_agent_app_debug_conversation( - app_id=backing_app_id, - account_id=account_id, - ) - previous_conversation: tuple[str, str] | None = None - mapping = self._session.scalar( - select(AgentDebugConversation).where( - AgentDebugConversation.tenant_id == tenant_id, - AgentDebugConversation.agent_id == agent_id, - AgentDebugConversation.account_id == account_id, - AgentDebugConversation.draft_type == draft_type, + retired_binding_id: str | None = None + try: + conversation_id = self._create_agent_app_debug_conversation( + app_id=backing_app_id, + account_id=account_id, ) - ) - if mapping is None: - self._session.add( - AgentDebugConversation( - tenant_id=tenant_id, - agent_id=agent_id, - app_id=backing_app_id, - account_id=account_id, - draft_type=draft_type, - conversation_id=conversation_id, + mapping = self._session.scalar( + select(AgentDebugConversation).where( + AgentDebugConversation.tenant_id == tenant_id, + AgentDebugConversation.agent_id == agent_id, + AgentDebugConversation.account_id == account_id, + AgentDebugConversation.draft_type == AgentConfigDraftType.DRAFT, ) ) - else: - previous_app_id = mapping.app_id - previous_conversation_id = mapping.conversation_id - if previous_conversation_id: - previous_conversation = (previous_app_id or backing_app_id, previous_conversation_id) - mapping.app_id = backing_app_id - mapping.conversation_id = conversation_id - self._session.flush() - self._session.commit() - - if previous_conversation: - previous_app_id, previous_conversation_id = previous_conversation - self._cleanup_debug_conversation_runtime_sessions( + if mapping is None: + self._session.add( + AgentDebugConversation( + tenant_id=tenant_id, + agent_id=agent_id, + app_id=backing_app_id, + account_id=account_id, + draft_type=AgentConfigDraftType.DRAFT, + conversation_id=conversation_id, + ) + ) + else: + previous_app_id = mapping.app_id or backing_app_id + previous_conversation_id = mapping.conversation_id + if previous_conversation_id: + previous_conversation = self._session.scalar( + select(Conversation).where( + Conversation.id == previous_conversation_id, + Conversation.app_id == previous_app_id, + Conversation.from_source == ConversationFromSource.CONSOLE, + Conversation.from_account_id == account_id, + Conversation.is_deleted.is_(False), + ) + ) + if ( + previous_conversation is not None + and previous_conversation.agent_workspace_binding_id is not None + ): + binding_id = previous_conversation.agent_workspace_binding_id + binding = AgentWorkspaceService.get_active_binding( + session=self._session, + tenant_id=tenant_id, + binding_id=binding_id, + expected_owner_scope=WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=previous_app_id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=previous_conversation.id, + ), + ) + if binding is None or binding.agent_id != agent_id: + raise AgentWorkspaceNotFoundError( + "Agent debug Conversation participant Binding is unavailable" + ) + retired_binding_id = AgentWorkspaceService.retire_binding( + session=self._session, + tenant_id=tenant_id, + binding_id=binding_id, + ) + if retired_binding_id is None: + raise AgentWorkspaceNotFoundError( + "Agent debug Conversation participant Binding is unavailable" + ) + mapping.app_id = backing_app_id + mapping.conversation_id = conversation_id + self._session.flush() + self._session.commit() + except Exception: + self._session.rollback() + raise + if retired_binding_id is not None: + enqueue_agent_resource_collection( tenant_id=tenant_id, - agent_id=agent_id, - account_id=account_id, - draft_type=draft_type, - app_id=previous_app_id, - conversation_id=previous_conversation_id, + binding_ids=(retired_binding_id,), ) return conversation_id - def _cleanup_debug_conversation_runtime_sessions( - self, - *, - tenant_id: str, - agent_id: str, - account_id: str, - draft_type: AgentConfigDraftType, - app_id: str, - conversation_id: str, - ) -> None: + def reset_build_conversation(self, *, tenant_id: str, agent_id: str, account_id: str) -> str: + """Reset Build and retire its exact DEBUG_BUILD Draft-owned Binding. + + The mapping update, exact BUILD_DRAFT Binding retirement, and Draft + pointer clear commit in one transaction. Validation failures fail fast; + collection is enqueued only after commit. + """ + + agent = self._session.scalar( + select(Agent).where( + Agent.tenant_id == tenant_id, + Agent.id == agent_id, + Agent.status == AgentStatus.ACTIVE, + ) + ) + if agent is None: + raise AgentNotFoundError() + backing_app_id = self._ensure_workflow_agent_backing_app( + agent=agent, + account_id=agent.updated_by or agent.created_by, + ) + if not backing_app_id: + raise AgentNotFoundError() + + retired_binding_id: str | None = None try: - session_store = AgentAppRuntimeSessionStore() - stored_sessions = session_store.list_active_sessions_for_conversation( - tenant_id=tenant_id, - app_id=app_id, - conversation_id=conversation_id, + conversation_id = self._create_agent_app_debug_conversation( + app_id=backing_app_id, + account_id=account_id, ) - except Exception: - logger.warning( - "Failed to load Agent App runtime sessions for debug conversation refresh: " - "tenant_id=%s app_id=%s conversation_id=%s", - tenant_id, - app_id, - conversation_id, - exc_info=True, - ) - return - - for stored_session in stored_sessions: - try: - if stored_session.runtime_layer_specs: - payload = AgentBackendSessionCleanupPayload( - session_snapshot=stored_session.session_snapshot, - runtime_layer_specs=stored_session.runtime_layer_specs, - idempotency_key=( - f"{tenant_id}:{agent_id}:{account_id}:{draft_type.value}:{conversation_id}:" - "debug-session-cleanup:" - f"{stored_session.scope.agent_id}:" - f"{stored_session.scope.agent_config_snapshot_id or 'no-config'}:" - f"{stored_session.backend_run_id or 'no-run'}" - ), - metadata={ - "tenant_id": stored_session.scope.tenant_id, - "app_id": stored_session.scope.app_id, - "conversation_id": stored_session.scope.conversation_id, - "agent_id": stored_session.scope.agent_id, - "agent_config_snapshot_id": stored_session.scope.agent_config_snapshot_id, - "draft_type": draft_type.value, - "previous_agent_backend_run_id": stored_session.backend_run_id, - }, - ) - cleanup_conversation_agent_runtime_session.delay(payload.model_dump(mode="json")) - except Exception: - logger.warning( - "Failed to enqueue Agent backend cleanup for debug conversation refresh: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", - stored_session.scope.tenant_id, - stored_session.scope.app_id, - stored_session.scope.conversation_id, - stored_session.scope.agent_id, - stored_session.backend_run_id, - exc_info=True, + mapping = self._session.scalar( + select(AgentDebugConversation).where( + AgentDebugConversation.tenant_id == tenant_id, + AgentDebugConversation.agent_id == agent_id, + AgentDebugConversation.account_id == account_id, + AgentDebugConversation.draft_type == AgentConfigDraftType.DEBUG_BUILD, ) - finally: - try: - session_store.mark_cleaned( - scope=stored_session.scope, - backend_run_id=stored_session.backend_run_id, - ) - except Exception: - logger.warning( - "Failed to retire Agent App runtime session for debug conversation refresh: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", - stored_session.scope.tenant_id, - stored_session.scope.app_id, - stored_session.scope.conversation_id, - stored_session.scope.agent_id, - stored_session.backend_run_id, - exc_info=True, - ) + ) + build_draft = self._session.scalar( + select(AgentConfigDraft) + .where( + AgentConfigDraft.tenant_id == tenant_id, + AgentConfigDraft.agent_id == agent_id, + AgentConfigDraft.draft_type == AgentConfigDraftType.DEBUG_BUILD, + AgentConfigDraft.account_id == account_id, + ) + .order_by(AgentConfigDraft.updated_at.desc()) + .limit(1) + ) + if build_draft is not None and build_draft.agent_workspace_binding_id is not None: + binding_id = build_draft.agent_workspace_binding_id + binding = AgentWorkspaceService.get_active_binding( + session=self._session, + tenant_id=tenant_id, + binding_id=binding_id, + expected_owner_scope=WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=backing_app_id, + owner_type=AgentWorkspaceOwnerType.BUILD_DRAFT, + owner_id=build_draft.id, + ), + ) + if binding is None or binding.agent_id != agent_id: + raise AgentBuildSandboxNotFoundError() + retired_binding_id = AgentWorkspaceService.retire_binding( + session=self._session, + tenant_id=tenant_id, + binding_id=binding_id, + ) + if retired_binding_id is None: + raise AgentBuildSandboxNotFoundError() + build_draft.agent_workspace_binding_id = None - def load_or_create_agent_app_debug_conversation_ids_by_agent_id( + if mapping is None: + self._session.add( + AgentDebugConversation( + tenant_id=tenant_id, + agent_id=agent_id, + app_id=backing_app_id, + account_id=account_id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + conversation_id=conversation_id, + ) + ) + else: + mapping.app_id = backing_app_id + mapping.conversation_id = conversation_id + self._session.flush() + self._session.commit() + except Exception: + self._session.rollback() + raise + if retired_binding_id is not None: + enqueue_agent_resource_collection( + tenant_id=tenant_id, + binding_ids=(retired_binding_id,), + ) + return conversation_id + + def load_or_create_build_conversation_ids_by_agent_id( self, *, tenant_id: str, agents: list[Agent], account_id: str, - draft_type: AgentConfigDraftType = AgentConfigDraftType.DEBUG_BUILD, ) -> dict[str, str]: - """Return per-account scoped conversations for a page of Agent Apps.""" + """Return per-account Build conversations for a page of Agent Apps.""" conversation_ids_by_agent_id: dict[str, str] = {} changed = False @@ -839,7 +900,7 @@ class AgentRosterService: conversation_ids_by_agent_id[agent.id] = self._get_or_create_agent_app_debug_conversation( agent=agent, account_id=account_id, - draft_type=draft_type, + draft_type=AgentConfigDraftType.DEBUG_BUILD, ) changed = True if changed: @@ -1057,7 +1118,6 @@ class AgentRosterService: account_id=account.id, ) self._session.commit() - if FeatureService.get_system_features().webapp_auth.enabled: try: original_settings = EnterpriseService.WebAppAuth.get_app_access_mode_by_id(source_app.id) @@ -1211,13 +1271,38 @@ class AgentRosterService: def archive_roster_agent(self, *, tenant_id: str, agent_id: str, account_id: str) -> None: agent = self._get_agent(tenant_id=tenant_id, agent_id=agent_id, roster_only=True) - if agent.status == AgentStatus.ARCHIVED: - return - agent.status = AgentStatus.ARCHIVED - agent.archived_by = account_id - agent.archived_at = naive_utc_now() - agent.updated_by = account_id + retired_binding_ids: list[str] = [] + if agent.status != AgentStatus.ARCHIVED: + agent.status = AgentStatus.ARCHIVED + agent.archived_by = account_id + agent.archived_at = naive_utc_now() + agent.updated_by = account_id + bindings = self._session.scalars( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.agent_id == agent_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + ).all() + for binding in bindings: + retired_id = AgentWorkspaceService.retire_binding( + session=self._session, + tenant_id=tenant_id, + binding_id=binding.id, + ) + if retired_id is not None: + retired_binding_ids.append(retired_id) + retired_snapshot_ids = AgentHomeSnapshotService.retire_all_for_agent( + session=self._session, + tenant_id=tenant_id, + agent_id=agent_id, + ) self._session.commit() + enqueue_agent_resource_collection( + tenant_id=tenant_id, + binding_ids=retired_binding_ids, + home_snapshot_ids=retired_snapshot_ids, + ) @staticmethod def _visible_version_operations(agent: Agent) -> set[AgentConfigRevisionOperation]: @@ -1382,9 +1467,11 @@ class AgentRosterService: account_id=None, draft_owner_key="", created_by=account_id, + home_snapshot_id=version.home_snapshot_id, ) self._session.add(draft) draft.base_snapshot_id = version.id + draft.home_snapshot_id = version.home_snapshot_id draft.config_snapshot = AgentSoulConfig.model_validate(version.config_snapshot_dict) draft.updated_by = account_id agent.active_config_is_published = version.id == agent.active_config_snapshot_id diff --git a/api/services/agent/workflow_publish_service.py b/api/services/agent/workflow_publish_service.py index e74a93a899b..b43f3091f88 100644 --- a/api/services/agent/workflow_publish_service.py +++ b/api/services/agent/workflow_publish_service.py @@ -24,6 +24,7 @@ from models.agent_config_entities import ( WorkflowNodeJobConfig, WorkflowPreviousNodeOutputRef, ) +from models.model import App from models.workflow import Workflow from services.agent.composer_validator import ComposerConfigValidator from services.agent.prompt_mentions import ( @@ -224,7 +225,7 @@ class WorkflowAgentPublishService: session: Session, draft_workflow: Workflow, account_id: str, - ) -> None: + ) -> set[str]: agent_nodes = dict(WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict)) existing_bindings = list( session.scalars( @@ -237,9 +238,12 @@ class WorkflowAgentPublishService: ).all() ) existing_by_node_id = {binding.node_id: binding for binding in existing_bindings} + retirement_candidates: set[str] = set() for binding in existing_bindings: if binding.node_id not in agent_nodes: + if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id: + retirement_candidates.add(binding.agent_id) session.delete(binding) for node_id, node_data in agent_nodes.items(): @@ -252,16 +256,34 @@ class WorkflowAgentPublishService: not binding_payload.get("agent_id") or not binding_payload.get("current_snapshot_id") ): continue + existing_binding = existing_by_node_id.get(node_id) + replaced_inline_agent_id = ( + existing_binding.agent_id + if existing_binding is not None + and existing_binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT + and existing_binding.agent_id + else None + ) cls._sync_agent_binding_for_node( session=session, draft_workflow=draft_workflow, node_id=node_id, node_data=node_data, node_binding=binding_payload, - existing_binding=existing_by_node_id.get(node_id), + existing_binding=existing_binding, account_id=account_id, ) + if ( + replaced_inline_agent_id + and existing_binding is not None + and ( + existing_binding.binding_type != WorkflowAgentBindingType.INLINE_AGENT + or existing_binding.agent_id != replaced_inline_agent_id + ) + ): + retirement_candidates.add(replaced_inline_agent_id) session.flush() + return retirement_candidates @classmethod def sync_roster_agent_bindings_for_draft( @@ -270,8 +292,8 @@ class WorkflowAgentPublishService: session: Session, draft_workflow: Workflow, account_id: str, - ) -> None: - cls.sync_agent_bindings_for_draft( + ) -> set[str]: + return cls.sync_agent_bindings_for_draft( session=session, draft_workflow=draft_workflow, account_id=account_id, @@ -561,12 +583,32 @@ class WorkflowAgentPublishService: session: Session, draft_workflow: Workflow, published_workflow: Workflow, - ) -> None: + ) -> set[str]: + current_workflow_id = session.scalar( + select(App.workflow_id).where( + App.tenant_id == draft_workflow.tenant_id, + App.id == draft_workflow.app_id, + ) + ) + retirement_candidates: set[str] = set() + if current_workflow_id: + retirement_candidates = { + agent_id + for agent_id in session.scalars( + select(WorkflowAgentNodeBinding.agent_id).where( + WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id, + WorkflowAgentNodeBinding.app_id == draft_workflow.app_id, + WorkflowAgentNodeBinding.workflow_id == current_workflow_id, + WorkflowAgentNodeBinding.binding_type == WorkflowAgentBindingType.INLINE_AGENT, + ) + ).all() + if agent_id + } node_ids = { node_id for node_id, _node_data in WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict) } if not node_ids: - return + return retirement_candidates bindings = session.scalars( select(WorkflowAgentNodeBinding).where( @@ -578,7 +620,7 @@ class WorkflowAgentPublishService: ) ).all() if not bindings: - return + return retirement_candidates agents_by_id = { agent.id: agent @@ -611,6 +653,7 @@ class WorkflowAgentPublishService: updated_by=binding.updated_by, ) session.add(copied) + return retirement_candidates @classmethod def restore_agent_node_bindings_to_draft( @@ -620,7 +663,7 @@ class WorkflowAgentPublishService: source_workflow: Workflow, draft_workflow: Workflow, account_id: str, - ) -> None: + ) -> set[str]: """Replace draft bindings with the frozen bindings of a published workflow.""" existing = session.scalars( @@ -631,6 +674,11 @@ class WorkflowAgentPublishService: WorkflowAgentNodeBinding.workflow_version == cls._DRAFT_WORKFLOW_VERSION, ) ).all() + retirement_candidates = { + binding.agent_id + for binding in existing + if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id + } for binding in existing: session.delete(binding) @@ -681,3 +729,4 @@ class WorkflowAgentPublishService: ) ) session.flush() + return retirement_candidates diff --git a/api/services/agent/workspace_service.py b/api/services/agent/workspace_service.py new file mode 100644 index 00000000000..f57ee7655b6 --- /dev/null +++ b/api/services/agent/workspace_service.py @@ -0,0 +1,469 @@ +"""Own Workspace and AgentWorkspaceBinding product lifecycle. + +Dify API is the lifecycle ledger. Dify Agent only executes physical create, +acquire, and destroy operations selected by this service. Retire methods only +mutate the caller's transaction; collection performs network I/O after commit +and deletes ledger rows only after idempotent physical cleanup succeeds. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +from dify_agent.client import Client +from dify_agent.protocol import CreateExecutionBindingRequest, DestroyExecutionBindingRequest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from configs import dify_config +from core.db.session_factory import session_factory +from libs.datetime_utils import naive_utc_now +from libs.uuid_utils import uuidv7 +from models.agent import ( + AgentConfigVersionKind, + AgentHomeSnapshot, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) + +logger = logging.getLogger(__name__) + + +class AgentWorkspaceError(RuntimeError): + pass + + +class AgentWorkspaceNotFoundError(AgentWorkspaceError): + pass + + +class AgentWorkspaceBindingGenerationMismatchError(AgentWorkspaceError): + pass + + +@dataclass(frozen=True, slots=True) +class WorkspaceOwnerScope: + tenant_id: str + app_id: str + owner_type: AgentWorkspaceOwnerType + owner_id: str + owner_scope_key: str = "root" + + +class AgentWorkspaceService: + """Allocate and manage working-environment resources. + + A Binding ID is the participant identity. Product callers persist that ID + and use :meth:`get_active_binding`; Agent and Workspace attributes are not + participant lookup keys. + """ + + @classmethod + def resolve_active_workspace(cls, *, session: Session, scope: WorkspaceOwnerScope) -> AgentWorkspace | None: + return session.scalar( + select(AgentWorkspace).where( + AgentWorkspace.tenant_id == scope.tenant_id, + AgentWorkspace.app_id == scope.app_id, + AgentWorkspace.owner_type == scope.owner_type, + AgentWorkspace.owner_id == scope.owner_id, + AgentWorkspace.owner_scope_key == scope.owner_scope_key, + AgentWorkspace.status == AgentWorkingResourceStatus.ACTIVE, + ) + ) + + @classmethod + def get_active_binding( + cls, + *, + session: Session, + tenant_id: str, + binding_id: str, + expected_owner_scope: WorkspaceOwnerScope, + ) -> AgentWorkspaceBinding | None: + return session.scalar( + select(AgentWorkspaceBinding) + .join( + AgentWorkspace, + (AgentWorkspace.tenant_id == AgentWorkspaceBinding.tenant_id) + & (AgentWorkspace.id == AgentWorkspaceBinding.workspace_id), + ) + .where( + AgentWorkspaceBinding.id == binding_id, + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + AgentWorkspace.tenant_id == expected_owner_scope.tenant_id, + AgentWorkspace.app_id == expected_owner_scope.app_id, + AgentWorkspace.owner_type == expected_owner_scope.owner_type, + AgentWorkspace.owner_id == expected_owner_scope.owner_id, + AgentWorkspace.owner_scope_key == expected_owner_scope.owner_scope_key, + AgentWorkspace.status == AgentWorkingResourceStatus.ACTIVE, + ) + ) + + @classmethod + def create_binding( + cls, + *, + session: Session, + scope: WorkspaceOwnerScope, + agent_id: str, + base_home_snapshot_id: str, + agent_config_version_id: str, + agent_config_version_kind: AgentConfigVersionKind, + ) -> AgentWorkspaceBinding: + """Allocate one new participant in the caller-owned transaction. + + After backend creation returns successfully, any later Python, flush, + or commit failure may leave an orphan. Dify API does not perform + cross-system compensation; a future global reconciler is responsible + for those orphans. Backend-local cleanup applies only when creation + fails before the backend returns success. + """ + + home_snapshot = session.scalar( + select(AgentHomeSnapshot).where( + AgentHomeSnapshot.id == base_home_snapshot_id, + AgentHomeSnapshot.tenant_id == scope.tenant_id, + AgentHomeSnapshot.agent_id == agent_id, + AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE, + ) + ) + if home_snapshot is None: + raise AgentWorkspaceNotFoundError("base Home Snapshot is unavailable") + workspace = cls.resolve_active_workspace(session=session, scope=scope) + workspace_id = workspace.id if workspace is not None else str(uuidv7()) + binding_id = str(uuidv7()) + with cls._client() as client: + allocation = client.create_execution_binding_sync( + CreateExecutionBindingRequest( + tenant_id=scope.tenant_id, + agent_id=agent_id, + binding_id=binding_id, + workspace_id=workspace_id, + existing_workspace_ref=workspace.backend_workspace_ref if workspace is not None else None, + home_snapshot_ref=home_snapshot.snapshot_ref, + ) + ) + if workspace is not None and allocation.workspace_ref != workspace.backend_workspace_ref: + raise AgentWorkspaceError("backend changed the existing Workspace ref") + if workspace is None: + workspace = AgentWorkspace( + id=workspace_id, + tenant_id=scope.tenant_id, + app_id=scope.app_id, + owner_type=scope.owner_type, + owner_id=scope.owner_id, + owner_scope_key=scope.owner_scope_key, + backend_workspace_ref=allocation.workspace_ref, + status=AgentWorkingResourceStatus.ACTIVE, + active_guard=1, + ) + session.add(workspace) + binding = AgentWorkspaceBinding( + id=binding_id, + tenant_id=scope.tenant_id, + app_id=scope.app_id, + workspace_id=workspace_id, + agent_id=agent_id, + base_home_snapshot_id=base_home_snapshot_id, + agent_config_version_id=agent_config_version_id, + agent_config_version_kind=agent_config_version_kind, + backend_binding_ref=allocation.binding_ref, + status=AgentWorkingResourceStatus.ACTIVE, + ) + session.add(binding) + return binding + + @classmethod + def save_binding_session_snapshot( + cls, + *, + tenant_id: str, + binding_id: str, + session_snapshot: str, + pending_form_id: str | None = None, + pending_tool_call_id: str | None = None, + ) -> None: + with session_factory.create_session() as session: + binding = session.scalar( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.id == binding_id, + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + ) + if binding is None: + raise AgentWorkspaceNotFoundError("ACTIVE Binding is unavailable") + binding.session_snapshot = session_snapshot + binding.pending_form_id = pending_form_id + binding.pending_tool_call_id = pending_tool_call_id + session.commit() + + @classmethod + def retire_binding(cls, *, session: Session, tenant_id: str, binding_id: str) -> str | None: + binding = session.scalar( + select(AgentWorkspaceBinding) + .where( + AgentWorkspaceBinding.id == binding_id, + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + .with_for_update() + ) + if binding is None: + return None + workspace = session.scalar( + select(AgentWorkspace) + .where( + AgentWorkspace.id == binding.workspace_id, + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.status == AgentWorkingResourceStatus.ACTIVE, + ) + .with_for_update() + ) + now = naive_utc_now() + binding.status = AgentWorkingResourceStatus.RETIRED + binding.retired_at = now + if workspace is not None: + other_binding = session.scalar( + select(AgentWorkspaceBinding.id).where( + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.workspace_id == workspace.id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + AgentWorkspaceBinding.id != binding.id, + ) + ) + if other_binding is None: + workspace.status = AgentWorkingResourceStatus.RETIRED + workspace.active_guard = None + workspace.retired_at = now + return binding.id + + @classmethod + def retire_workspace(cls, *, session: Session, tenant_id: str, workspace_id: str) -> str | None: + workspace = session.scalar( + select(AgentWorkspace) + .where( + AgentWorkspace.id == workspace_id, + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.status == AgentWorkingResourceStatus.ACTIVE, + ) + .with_for_update() + ) + if workspace is None: + return None + now = naive_utc_now() + workspace.status = AgentWorkingResourceStatus.RETIRED + workspace.active_guard = None + workspace.retired_at = now + bindings = session.scalars( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.workspace_id == workspace.id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + ).all() + for binding in bindings: + binding.status = AgentWorkingResourceStatus.RETIRED + binding.retired_at = now + return workspace.id + + @classmethod + def retire_all_for_app(cls, *, session: Session, tenant_id: str, app_id: str) -> list[str]: + """Retire all ACTIVE Workspaces owned by an App in the caller's transaction.""" + + workspaces = session.scalars( + select(AgentWorkspace).where( + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.app_id == app_id, + AgentWorkspace.status == AgentWorkingResourceStatus.ACTIVE, + ) + ).all() + retired: list[str] = [] + for workspace in workspaces: + workspace_id = cls.retire_workspace( + session=session, + tenant_id=tenant_id, + workspace_id=workspace.id, + ) + if workspace_id is not None: + retired.append(workspace_id) + return retired + + @classmethod + def collect_retired_binding(cls, *, tenant_id: str, binding_id: str) -> None: + try: + cls._collect_retired_binding(tenant_id=tenant_id, binding_id=binding_id) + except Exception: + logger.exception( + "Failed to collect retired Agent Workspace Binding", + extra={"tenant_id": tenant_id, "binding_id": binding_id}, + ) + + @classmethod + def _collect_retired_binding(cls, *, tenant_id: str, binding_id: str) -> None: + with session_factory.create_session() as session: + binding = session.scalar( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.id == binding_id, + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, + ) + ) + if binding is None: + return + backend_binding_ref = binding.backend_binding_ref + workspace = session.scalar( + select(AgentWorkspace).where( + AgentWorkspace.id == binding.workspace_id, + AgentWorkspace.tenant_id == tenant_id, + ) + ) + if workspace is not None and workspace.status == AgentWorkingResourceStatus.RETIRED: + workspace_id = workspace.id + else: + workspace_id = None + if workspace_id is not None: + cls.collect_retired_workspace(tenant_id=tenant_id, workspace_id=workspace_id) + return + try: + with cls._client() as client: + client.destroy_execution_binding_sync( + DestroyExecutionBindingRequest( + binding_ref=backend_binding_ref, + destroy_workspace=False, + ) + ) + except Exception: + logger.exception( + "Failed to collect retired Agent Workspace Binding", + extra={"tenant_id": tenant_id, "binding_id": binding_id}, + ) + return + with session_factory.create_session() as session: + binding = session.scalar( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.id == binding_id, + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, + ) + ) + if binding is not None: + session.delete(binding) + session.commit() + + @classmethod + def collect_retired_workspace(cls, *, tenant_id: str, workspace_id: str) -> None: + try: + cls._collect_retired_workspace(tenant_id=tenant_id, workspace_id=workspace_id) + except Exception: + logger.exception( + "Failed to collect retired Agent Workspace", + extra={"tenant_id": tenant_id, "workspace_id": workspace_id}, + ) + + @classmethod + def _collect_retired_workspace(cls, *, tenant_id: str, workspace_id: str) -> None: + with session_factory.create_session() as session: + workspace = session.scalar( + select(AgentWorkspace).where( + AgentWorkspace.id == workspace_id, + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.status == AgentWorkingResourceStatus.RETIRED, + ) + ) + if workspace is None: + return + bindings = session.scalars( + select(AgentWorkspaceBinding) + .where( + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.workspace_id == workspace_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, + ) + .order_by(AgentWorkspaceBinding.created_at) + ).all() + if not bindings: + logger.error( + "RETIRED Workspace has no Binding available for physical collection", + extra={"tenant_id": tenant_id, "workspace_id": workspace_id}, + ) + return + anchor = bindings[0] + remaining_ids = [binding.id for binding in bindings[1:]] + workspace_ref = workspace.backend_workspace_ref + binding_ref = anchor.backend_binding_ref + anchor_id = anchor.id + try: + with cls._client() as client: + client.destroy_execution_binding_sync( + DestroyExecutionBindingRequest( + binding_ref=binding_ref, + workspace_ref=workspace_ref, + destroy_workspace=True, + ) + ) + except Exception: + logger.exception( + "Failed to collect retired Agent Workspace", + extra={"tenant_id": tenant_id, "workspace_id": workspace_id, "binding_id": anchor_id}, + ) + return + with session_factory.create_session() as session: + stored_workspace = session.scalar( + select(AgentWorkspace).where( + AgentWorkspace.id == workspace_id, + AgentWorkspace.tenant_id == tenant_id, + AgentWorkspace.status == AgentWorkingResourceStatus.RETIRED, + ) + ) + stored_anchor = session.scalar( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.id == anchor_id, + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, + ) + ) + if stored_workspace is not None: + session.delete(stored_workspace) + if stored_anchor is not None: + session.delete(stored_anchor) + session.commit() + for remaining_id in remaining_ids: + cls.collect_retired_binding(tenant_id=tenant_id, binding_id=remaining_id) + + @staticmethod + def validate_binding_generation( + binding: AgentWorkspaceBinding, + *, + base_home_snapshot_id: str, + agent_config_version_id: str, + agent_config_version_kind: AgentConfigVersionKind, + ) -> None: + if ( + binding.base_home_snapshot_id != base_home_snapshot_id + or binding.agent_config_version_id != agent_config_version_id + or binding.agent_config_version_kind != agent_config_version_kind + ): + raise AgentWorkspaceBindingGenerationMismatchError( + "ACTIVE Binding belongs to a different Agent config/Home generation" + ) + + @staticmethod + def _client() -> Client: + base_url = dify_config.AGENT_BACKEND_BASE_URL + if not base_url: + raise AgentWorkspaceError("Dify Agent backend is required for Workspace operations") + return Client(base_url=base_url) + + +__all__ = [ + "AgentWorkspaceBindingGenerationMismatchError", + "AgentWorkspaceError", + "AgentWorkspaceNotFoundError", + "AgentWorkspaceService", + "WorkspaceOwnerScope", +] diff --git a/api/services/agent_app_sandbox_service.py b/api/services/agent_app_sandbox_service.py index 3f5a0bf41b2..9661c30c7d3 100644 --- a/api/services/agent_app_sandbox_service.py +++ b/api/services/agent_app_sandbox_service.py @@ -1,39 +1,40 @@ -"""Resolve and proxy sandbox file access for Agent App and workflow Agent sessions. - -These services keep product-facing locators (conversation, workflow run, node) -on the API boundary and translate them into the agent backend's -``SandboxLocator`` using persisted non-sensitive runtime layer specs plus the -saved Agenton session snapshot. Upload responses stay console-facing here: the -agent backend still returns a canonical ToolFile mapping, while this API layer -re-resolves that mapping into a signed browser download URL. -""" +"""Resolve product locators to ACTIVE Workspace Bindings and proxy file access.""" from __future__ import annotations import urllib.parse from collections.abc import Callable -from typing import Any +from typing import Any, Literal, cast -from agenton.compositor import CompositorSessionSnapshot from dify_agent.client import Client -from dify_agent.protocol import RuntimeLayerSpec, SandboxLocator, build_sandbox_locator_from_layer_specs -from pydantic import BaseModel, TypeAdapter +from dify_agent.layers.execution_context import ( + DifyExecutionContextAgentConfigVersionKind, + DifyExecutionContextLayerConfig, +) +from dify_agent.protocol import WorkspaceListResponse, WorkspaceReadResponse, WorkspaceUploadRequest +from pydantic import BaseModel from sqlalchemy import select from sqlalchemy.orm import Session from configs import dify_config -from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore from core.app.file_access import DatabaseFileAccessController from core.app.workflow.file_runtime import DifyWorkflowFileRuntime +from core.db.session_factory import session_factory from factories import file_factory -from models.agent import AgentRuntimeSessionOwnerType, WorkflowAgentRuntimeSession, WorkflowAgentRuntimeSessionStatus - -_RUNTIME_LAYER_SPECS_ADAPTER: TypeAdapter[list[RuntimeLayerSpec]] = TypeAdapter(list[RuntimeLayerSpec]) +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigDraftType, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) +from models.model import App, Conversation +from models.workflow import WorkflowNodeExecutionModel +from services.agent.roster_service import AgentRosterService +from services.agent.workspace_service import AgentWorkspaceService, WorkspaceOwnerScope class AgentSandboxInspectorError(Exception): - """A sandbox inspection failure mapped to an HTTP status by the controller.""" - code: str message: str status_code: int @@ -46,82 +47,200 @@ class AgentSandboxInspectorError(Exception): class AgentSandboxInfo(BaseModel): - """Basic Agent App sandbox metadata returned after a successful availability probe.""" - - session_id: str workspace_cwd: str class AgentSandboxUploadDownload(BaseModel): - """Signed browser download URL for one sandbox upload result.""" - url: str class AgentAppSandboxService: - """Inspect and proxy file access for an Agent App conversation sandbox.""" - - def __init__( - self, - *, - session_store: AgentAppRuntimeSessionStore | None = None, - client_factory: Callable[[], Client] | None = None, - ) -> None: - self._session_store = session_store or AgentAppRuntimeSessionStore() + def __init__(self, *, client_factory: Callable[[], Client] | None = None) -> None: self._client_factory = client_factory or _default_client_factory - def get_info(self, *, tenant_id: str, app_id: str, conversation_id: str) -> AgentSandboxInfo: - locator = self._resolve_locator(tenant_id=tenant_id, app_id=app_id, conversation_id=conversation_id) - session_id, workspace_cwd = _extract_shell_workspace_or_raise( - snapshot=locator.session_snapshot, - not_found_message="this conversation's agent has no sandbox workspace", - ) - - return AgentSandboxInfo( - session_id=session_id, - workspace_cwd=workspace_cwd, - ) - - def list_files(self, *, tenant_id: str, app_id: str, conversation_id: str, path: str): - locator = self._resolve_locator(tenant_id=tenant_id, app_id=app_id, conversation_id=conversation_id) - return self._client_factory().list_sandbox_files_sync(locator, path) - - def read_file(self, *, tenant_id: str, app_id: str, conversation_id: str, path: str): - locator = self._resolve_locator(tenant_id=tenant_id, app_id=app_id, conversation_id=conversation_id) - return self._client_factory().read_sandbox_file_sync(locator, path) - - def upload_file( - self, *, tenant_id: str, app_id: str, conversation_id: str, path: str - ) -> AgentSandboxUploadDownload: - locator = self._resolve_locator(tenant_id=tenant_id, app_id=app_id, conversation_id=conversation_id) - uploaded = self._client_factory().upload_sandbox_file_sync(locator, path) - return _upload_download_response( - tenant_id=tenant_id, - file_mapping=uploaded.file.model_dump(mode="python"), - ) - - def _resolve_locator(self, *, tenant_id: str, app_id: str, conversation_id: str) -> SandboxLocator: - stored = self._session_store.load_active_session_for_conversation( + def get_info( + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: Literal["conversation", "build_draft"], + caller_id: str, + account_id: str, + ) -> AgentSandboxInfo: + self._resolve_binding( tenant_id=tenant_id, app_id=app_id, - conversation_id=conversation_id, + agent_id=agent_id, + caller_type=caller_type, + caller_id=caller_id, + account_id=account_id, ) - if stored is None: - raise AgentSandboxInspectorError( - "no_active_session", - "this conversation has no active sandbox session yet", - status_code=404, + return AgentSandboxInfo(workspace_cwd=".") + + def list_files( + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: Literal["conversation", "build_draft"], + caller_id: str, + account_id: str, + path: str, + ) -> WorkspaceListResponse: + binding = self._resolve_binding( + tenant_id=tenant_id, + app_id=app_id, + agent_id=agent_id, + caller_type=caller_type, + caller_id=caller_id, + account_id=account_id, + ) + with self._client_factory() as client: + return client.list_workspace_files_sync(binding.backend_binding_ref, path) + + def read_file( + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: Literal["conversation", "build_draft"], + caller_id: str, + account_id: str, + path: str, + ) -> WorkspaceReadResponse: + binding = self._resolve_binding( + tenant_id=tenant_id, + app_id=app_id, + agent_id=agent_id, + caller_type=caller_type, + caller_id=caller_id, + account_id=account_id, + ) + with self._client_factory() as client: + return client.read_workspace_file_sync(binding.backend_binding_ref, path) + + def upload_file( + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: Literal["conversation", "build_draft"], + caller_id: str, + account_id: str, + path: str, + ) -> AgentSandboxUploadDownload: + binding = self._resolve_binding( + tenant_id=tenant_id, + app_id=app_id, + agent_id=agent_id, + caller_type=caller_type, + caller_id=caller_id, + account_id=account_id, + ) + with self._client_factory() as client: + uploaded = client.upload_workspace_file_sync( + WorkspaceUploadRequest( + backend_binding_ref=binding.backend_binding_ref, + path=path, + execution_context=DifyExecutionContextLayerConfig( + tenant_id=tenant_id, + app_id=app_id, + conversation_id=caller_id if caller_type == "conversation" else None, + agent_id=agent_id, + agent_config_version_id=binding.agent_config_version_id, + agent_config_version_kind=cast( + DifyExecutionContextAgentConfigVersionKind, + binding.agent_config_version_kind.value, + ), + agent_mode="agent_app", + invoke_from="debugger", + ), + ) ) - return _build_locator_or_raise( - snapshot=stored.session_snapshot, - runtime_layer_specs=stored.runtime_layer_specs, - not_found_message="this conversation's agent has no sandbox workspace", - ) + return _upload_download_response(tenant_id=tenant_id, file_mapping=uploaded.file.model_dump(mode="python")) + + @staticmethod + def _resolve_binding( + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: Literal["conversation", "build_draft"], + caller_id: str, + account_id: str, + ) -> AgentWorkspaceBinding: + with session_factory.create_session() as session: + caller: AgentConfigDraft | Conversation | None + if caller_type == "build_draft": + agent = session.scalar( + select(Agent).where( + Agent.id == agent_id, + Agent.tenant_id == tenant_id, + ) + ) + if agent is None or AgentRosterService.runtime_backing_app_id(agent) != app_id: + caller = None + else: + caller = session.scalar( + select(AgentConfigDraft).where( + AgentConfigDraft.id == caller_id, + AgentConfigDraft.tenant_id == tenant_id, + AgentConfigDraft.agent_id == agent_id, + AgentConfigDraft.account_id == account_id, + AgentConfigDraft.draft_type == AgentConfigDraftType.DEBUG_BUILD, + ) + ) + owner_scope = WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.BUILD_DRAFT, + owner_id=caller_id, + ) + else: + caller = session.scalar( + select(Conversation) + .join(App, App.id == Conversation.app_id) + .where( + App.tenant_id == tenant_id, + Conversation.app_id == app_id, + Conversation.id == caller_id, + Conversation.from_account_id == account_id, + Conversation.is_deleted.is_(False), + ) + ) + owner_scope = WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=caller_id, + ) + if caller is None or caller.agent_workspace_binding_id is None: + raise AgentSandboxInspectorError( + "no_active_binding", + "this caller has no active Agent Workspace Binding", + status_code=404, + ) + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=tenant_id, + binding_id=caller.agent_workspace_binding_id, + expected_owner_scope=owner_scope, + ) + if binding is None or binding.agent_id != agent_id: + raise AgentSandboxInspectorError( + "no_active_binding", + "this caller has no active Agent Workspace Binding", + status_code=404, + ) + session.expunge(binding) + return binding class WorkflowAgentSandboxService: - """List/read/upload files in a workflow Agent node sandbox.""" - def __init__(self, *, client_factory: Callable[[], Client] | None = None) -> None: self._client_factory = client_factory or _default_client_factory @@ -132,11 +251,11 @@ class WorkflowAgentSandboxService: app_id: str, workflow_run_id: str, node_id: str, - node_execution_id: str | None, + node_execution_id: str, path: str, session: Session, - ): - locator = self._resolve_locator( + ) -> WorkspaceListResponse: + binding = self._resolve_binding( tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run_id, @@ -144,7 +263,8 @@ class WorkflowAgentSandboxService: node_execution_id=node_execution_id, session=session, ) - return self._client_factory().list_sandbox_files_sync(locator, path) + with self._client_factory() as client: + return client.list_workspace_files_sync(binding.backend_binding_ref, path) def read_file( self, @@ -153,11 +273,11 @@ class WorkflowAgentSandboxService: app_id: str, workflow_run_id: str, node_id: str, - node_execution_id: str | None, + node_execution_id: str, path: str, session: Session, - ): - locator = self._resolve_locator( + ) -> WorkspaceReadResponse: + binding = self._resolve_binding( tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run_id, @@ -165,7 +285,8 @@ class WorkflowAgentSandboxService: node_execution_id=node_execution_id, session=session, ) - return self._client_factory().read_sandbox_file_sync(locator, path) + with self._client_factory() as client: + return client.read_workspace_file_sync(binding.backend_binding_ref, path) def upload_file( self, @@ -174,11 +295,11 @@ class WorkflowAgentSandboxService: app_id: str, workflow_run_id: str, node_id: str, - node_execution_id: str | None, + node_execution_id: str, path: str, session: Session, ) -> AgentSandboxUploadDownload: - locator = self._resolve_locator( + binding = self._resolve_binding( tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run_id, @@ -186,118 +307,97 @@ class WorkflowAgentSandboxService: node_execution_id=node_execution_id, session=session, ) - uploaded = self._client_factory().upload_sandbox_file_sync(locator, path) - return _upload_download_response( - tenant_id=tenant_id, - file_mapping=uploaded.file.model_dump(mode="python"), - ) + with self._client_factory() as client: + uploaded = client.upload_workspace_file_sync( + WorkspaceUploadRequest( + backend_binding_ref=binding.backend_binding_ref, + path=path, + execution_context=DifyExecutionContextLayerConfig( + tenant_id=tenant_id, + app_id=app_id, + workflow_run_id=workflow_run_id, + node_id=node_id, + agent_id=binding.agent_id, + agent_config_version_id=binding.agent_config_version_id, + agent_config_version_kind=cast( + DifyExecutionContextAgentConfigVersionKind, + binding.agent_config_version_kind.value, + ), + agent_mode="workflow_run", + invoke_from="debugger", + ), + ) + ) + return _upload_download_response(tenant_id=tenant_id, file_mapping=uploaded.file.model_dump(mode="python")) - def _resolve_locator( - self, + @staticmethod + def _resolve_binding( *, tenant_id: str, app_id: str, workflow_run_id: str, node_id: str, - node_execution_id: str | None, + node_execution_id: str, session: Session, - ) -> SandboxLocator: - """Resolve one workflow Agent sandbox from product-facing identifiers. - - Callers may target either a specific node execution or the current node - as a whole. When ``node_execution_id`` is provided, lookup narrows to - that execution's ACTIVE runtime-session row. When it is omitted, the - service falls back to the most recently updated ACTIVE session for the - same ``workflow_run_id + node_id`` pair so console sandbox inspection can - still work from the broader workflow/node locator. - """ - stmt = select(WorkflowAgentRuntimeSession).where( - WorkflowAgentRuntimeSession.owner_type == AgentRuntimeSessionOwnerType.WORKFLOW_RUN, - WorkflowAgentRuntimeSession.tenant_id == tenant_id, - WorkflowAgentRuntimeSession.app_id == app_id, - WorkflowAgentRuntimeSession.workflow_run_id == workflow_run_id, - WorkflowAgentRuntimeSession.node_id == node_id, - WorkflowAgentRuntimeSession.status == WorkflowAgentRuntimeSessionStatus.ACTIVE, + ) -> AgentWorkspaceBinding: + execution = session.scalar( + select(WorkflowNodeExecutionModel).where( + WorkflowNodeExecutionModel.id == node_execution_id, + WorkflowNodeExecutionModel.tenant_id == tenant_id, + WorkflowNodeExecutionModel.app_id == app_id, + WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id, + WorkflowNodeExecutionModel.node_id == node_id, + ) ) - if node_execution_id: - stmt = stmt.where(WorkflowAgentRuntimeSession.node_execution_id == node_execution_id) - stmt = stmt.order_by(WorkflowAgentRuntimeSession.updated_at.desc()).limit(1) - - row = session.scalar(stmt) - - if row is None: + process_data = execution.process_data_dict if execution is not None else None + workflow_agent_binding_id = process_data.get("workflow_agent_binding_id") if process_data is not None else None + if ( + execution is None + or execution.agent_workspace_binding_id is None + or not isinstance(workflow_agent_binding_id, str) + ): raise AgentSandboxInspectorError( - "no_active_session", - "this workflow Agent node has no active sandbox session yet", + "no_active_binding", + "this Workflow Agent node execution has no active Workspace Binding", status_code=404, ) - return _build_locator_or_raise( - snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), - runtime_layer_specs=_deserialize_runtime_layer_specs(row.composition_layer_specs), - not_found_message="this workflow Agent node has no sandbox workspace", + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=tenant_id, + binding_id=execution.agent_workspace_binding_id, + expected_owner_scope=WorkspaceOwnerScope( + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN, + owner_id=workflow_run_id, + owner_scope_key=f"{node_id}:{workflow_agent_binding_id}", + ), ) - - -def _build_locator_or_raise( - *, - snapshot: CompositorSessionSnapshot, - runtime_layer_specs: list[RuntimeLayerSpec], - not_found_message: str, -) -> SandboxLocator: - try: - return build_sandbox_locator_from_layer_specs( - layer_specs=runtime_layer_specs, - session_snapshot=snapshot, - ) - except ValueError as exc: - raise AgentSandboxInspectorError("no_sandbox", not_found_message, status_code=404) from exc - - -def _extract_shell_workspace_or_raise( - *, - snapshot: CompositorSessionSnapshot, - not_found_message: str, -) -> tuple[str, str]: - shell_layer = next((layer for layer in snapshot.layers if layer.name == "shell"), None) - if shell_layer is None: - raise AgentSandboxInspectorError("no_sandbox", not_found_message, status_code=404) - - session_id = shell_layer.runtime_state.get("session_id") - workspace_cwd = shell_layer.runtime_state.get("workspace_cwd") - if not isinstance(session_id, str) or not isinstance(workspace_cwd, str): - raise AgentSandboxInspectorError("no_sandbox", not_found_message, status_code=404) - return session_id, workspace_cwd - - -def _deserialize_runtime_layer_specs(value: str | None) -> list[RuntimeLayerSpec]: - if not value: - return [] - return _RUNTIME_LAYER_SPECS_ADAPTER.validate_json(value) + if binding is None: + raise AgentSandboxInspectorError( + "no_active_binding", + "this Workflow Agent node execution has no active Workspace Binding", + status_code=404, + ) + return binding def _upload_download_response(*, tenant_id: str, file_mapping: dict[str, Any]) -> AgentSandboxUploadDownload: - """Resolve one uploaded ToolFile mapping into a signed external download URL.""" - controller = DatabaseFileAccessController() runtime = DifyWorkflowFileRuntime(file_access_controller=controller) try: - file = file_factory.build_from_mapping( - mapping=file_mapping, - tenant_id=tenant_id, - access_controller=controller, - ) + file = file_factory.build_from_mapping(mapping=file_mapping, tenant_id=tenant_id, access_controller=controller) url = runtime.resolve_file_url(file=file, for_external=True) except ValueError as exc: raise AgentSandboxInspectorError( - "sandbox_upload_download_unavailable", - "uploaded sandbox file could not be converted to a download URL", + "workspace_upload_download_unavailable", + "uploaded Workspace file could not be converted to a download URL", status_code=502, ) from exc - if not url: raise AgentSandboxInspectorError( - "sandbox_upload_download_unavailable", - "uploaded sandbox file does not support download URL generation", + "workspace_upload_download_unavailable", + "uploaded Workspace file does not support download URL generation", status_code=502, ) return AgentSandboxUploadDownload(url=_with_as_attachment(url)) @@ -315,7 +415,7 @@ def _default_client_factory() -> Client: if not base_url: raise AgentSandboxInspectorError( "inspector_unavailable", - "the sandbox file inspector is not available (agent backend not configured)", + "the Workspace file inspector is not available (Agent backend not configured)", status_code=503, ) return Client(base_url=base_url) diff --git a/api/services/app_dsl_service.py b/api/services/app_dsl_service.py index 8f41ca109fd..928eaa5b3d1 100644 --- a/api/services/app_dsl_service.py +++ b/api/services/app_dsl_service.py @@ -41,6 +41,7 @@ from models import Account, App, AppMode from models.model import AppModelConfig, AppModelConfigDict, IconType, load_annotation_reply_config from models.workflow import Workflow from services.agent.dsl_service import AgentDslService, AgentPackage +from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.workflow_publish_service import WorkflowAgentPublishService from services.dsl_content import DSL_MAX_SIZE, dsl_content_size from services.dsl_version import check_version_compatibility @@ -51,6 +52,7 @@ from services.errors.app import WorkflowNotFoundError from services.plugin.dependencies_analysis import DependenciesAnalysisService from services.workflow_draft_variable_service import WorkflowDraftVariableService from services.workflow_service import WorkflowService +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -546,7 +548,7 @@ class AppDslService: sync_agent_bindings=not raw_agent_packages, ) if raw_agent_packages: - _, warnings = AgentDslService(self._session).import_workflow_packages( + _, warnings, retirement_candidates = AgentDslService(self._session).import_workflow_packages( workflow=draft_workflow, portable_graph=graph, raw_packages=raw_agent_packages, @@ -557,6 +559,17 @@ class AppDslService: session=self._session, draft_workflow=draft_workflow, ) + self._session.commit() + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=app.tenant_id, + agent_ids=retirement_candidates, + account_id=account.id, + ) + enqueue_agent_resource_collection( + tenant_id=app.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) case AppMode.CHAT | AppMode.AGENT_CHAT | AppMode.COMPLETION: # Initialize model config model_config = data.get("model_config") diff --git a/api/services/app_service.py b/api/services/app_service.py index 6faa88eb114..de75b219ac2 100644 --- a/api/services/app_service.py +++ b/api/services/app_service.py @@ -26,17 +26,29 @@ from libs.datetime_utils import naive_utc_now from libs.login import current_user from libs.pagination import PaginatedResult, paginate_query from models import Account, AppStar -from models.agent import APP_BACKED_AGENT_SOURCES, Agent, AgentIconType, AgentScope, AgentStatus +from models.agent import ( + APP_BACKED_AGENT_SOURCES, + Agent, + AgentIconType, + AgentScope, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, +) from models.model import App, AppMode, AppModelConfig, IconType, Site, load_annotation_reply_config from models.tools import ApiToolProvider from models.workflow import Workflow from services.agent.errors import AgentNameConflictError +from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.retirement_service import WorkflowAgentRetirementService +from services.agent.workspace_service import AgentWorkspaceService from services.billing_service import BillingService from services.enterprise import rbac_service as enterprise_rbac_service from services.enterprise.enterprise_service import EnterpriseService from services.feature_service import FeatureService from services.openapi.visibility import apply_openapi_gate, is_openapi_visible from services.tag_service import TagService +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection from tasks.remove_app_and_related_data_task import remove_app_and_related_data_task logger = logging.getLogger(__name__) @@ -478,7 +490,14 @@ class AppService: session.delete(existing_star) - def create_app(self, tenant_id: str, params: CreateAppParams, account: Account, *, session: Session) -> App: + def create_app( + self, + tenant_id: str, + params: CreateAppParams, + account: Account, + *, + session: Session, + ) -> App: """ Create app :param tenant_id: tenant id @@ -943,18 +962,67 @@ class AppService: app_was_deleted.send(app) backing_agent = self._get_backing_agent_for_update(app, session=session) + workflow_agent_ids = session.scalars( + select(Agent.id).where( + Agent.tenant_id == app.tenant_id, + Agent.app_id == app.id, + Agent.scope == AgentScope.WORKFLOW_ONLY, + Agent.status == AgentStatus.ACTIVE, + ) + ).all() + account_id = current_user.id if current_user else None if backing_agent is not None: now = naive_utc_now() - account_id = getattr(current_user, "id", None) backing_agent.status = AgentStatus.ARCHIVED backing_agent.archived_by = account_id backing_agent.archived_at = now backing_agent.updated_by = account_id backing_agent.updated_at = now + retired_binding_ids: list[str] = [] + retired_snapshot_ids: list[str] = [] + if backing_agent is not None: + bindings = session.scalars( + select(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.tenant_id == app.tenant_id, + AgentWorkspaceBinding.agent_id == backing_agent.id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + ).all() + for binding in bindings: + binding_id = AgentWorkspaceService.retire_binding( + session=session, + tenant_id=app.tenant_id, + binding_id=binding.id, + ) + if binding_id is not None: + retired_binding_ids.append(binding_id) + retired_snapshot_ids = AgentHomeSnapshotService.retire_all_for_agent( + session=session, + tenant_id=app.tenant_id, + agent_id=backing_agent.id, + ) + + retired_workspace_ids = AgentWorkspaceService.retire_all_for_app( + session=session, + tenant_id=app.tenant_id, + app_id=app.id, + ) session.delete(app) session.commit() + workflow_binding_ids, workflow_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=app.tenant_id, + agent_ids=workflow_agent_ids, + account_id=account_id, + ) + enqueue_agent_resource_collection( + tenant_id=app.tenant_id, + workspace_ids=retired_workspace_ids, + binding_ids=[*retired_binding_ids, *workflow_binding_ids], + home_snapshot_ids=[*retired_snapshot_ids, *workflow_snapshot_ids], + ) + # clean up web app settings if FeatureService.get_system_features().webapp_auth.enabled: EnterpriseService.WebAppAuth.cleanup_webapp(app.id) diff --git a/api/services/conversation_service.py b/api/services/conversation_service.py index 6ab9c5a8a2b..3cd53f8249f 100644 --- a/api/services/conversation_service.py +++ b/api/services/conversation_service.py @@ -6,9 +6,7 @@ from typing import Any from sqlalchemy import asc, desc, func, or_, select from sqlalchemy.orm import Session -from clients.agent_backend import AgentBackendSessionCleanupPayload from configs import dify_config -from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore from core.app.entities.app_invoke_entities import InvokeFrom from core.llm_generator.llm_generator import LLMGenerator from factories import variable_factory @@ -16,7 +14,9 @@ from graphon.variables.types import SegmentType from libs.datetime_utils import naive_utc_now from libs.infinite_scroll_pagination import InfiniteScrollPagination from models import Account, ConversationVariable +from models.agent import AgentWorkspaceOwnerType from models.model import App, Conversation, EndUser, Message +from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope from services.errors.conversation import ( ConversationNotExistsError, ConversationVariableNotExistsError, @@ -24,7 +24,7 @@ from services.errors.conversation import ( LastConversationNotExistsError, ) from services.errors.message import MessageNotExistsError -from tasks.agent_backend_session_cleanup_task import cleanup_conversation_agent_runtime_session +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection from tasks.delete_conversation_task import delete_conversation_related_data logger = logging.getLogger(__name__) @@ -189,23 +189,30 @@ class ConversationService: """ Delete a conversation only if it belongs to the given user and app context. - Before removing the conversation row, this best-effort lifecycle path - enumerates any ACTIVE conversation-owned Agent backend runtime sessions, - enqueues asynchronous backend cleanup for rows with persisted runtime - layer specs, and then retires the local session rows even if enqueueing - fails. Conversation deletion and related-data cleanup scheduling still - proceed when that lifecycle bookkeeping only partially succeeds. + Conversation deletion is the product lifecycle boundary for its + Workspace. Physical collection happens only after the retire commit. Raises: ConversationNotExistsError: When the conversation is not visible to the current user. """ conversation = cls.get_conversation(app_model, conversation_id, user, session=session) - session_store = AgentAppRuntimeSessionStore() - stored_sessions = session_store.list_active_sessions_for_conversation( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - conversation_id=conversation.id, - ) + binding_id = conversation.agent_workspace_binding_id + retired_binding_id: str | None = None + if binding_id is not None: + owner_scope = WorkspaceOwnerScope( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=conversation.id, + ) + binding = AgentWorkspaceService.get_active_binding( + session=session, + tenant_id=app_model.tenant_id, + binding_id=binding_id, + expected_owner_scope=owner_scope, + ) + if binding is None: + raise AgentWorkspaceNotFoundError("Conversation participant Binding is unavailable") try: logger.info( @@ -213,66 +220,25 @@ class ConversationService: app_model.name, conversation_id, ) - for stored_session in stored_sessions: - try: - if stored_session.runtime_layer_specs: - payload = AgentBackendSessionCleanupPayload( - session_snapshot=stored_session.session_snapshot, - runtime_layer_specs=stored_session.runtime_layer_specs, - idempotency_key=( - f"{stored_session.scope.tenant_id}:{stored_session.scope.app_id}:" - f"{stored_session.scope.conversation_id}:agent-runtime-session-cleanup:" - f"{stored_session.scope.agent_id}:" - f"{stored_session.scope.agent_config_snapshot_id or 'no-config'}:" - f"{stored_session.backend_run_id or 'no-run'}" - ), - metadata={ - "tenant_id": stored_session.scope.tenant_id, - "app_id": stored_session.scope.app_id, - "conversation_id": stored_session.scope.conversation_id, - "agent_id": stored_session.scope.agent_id, - "agent_config_snapshot_id": stored_session.scope.agent_config_snapshot_id, - "previous_agent_backend_run_id": stored_session.backend_run_id, - }, - ) - cleanup_conversation_agent_runtime_session.delay(payload.model_dump(mode="json")) - except Exception: - logger.warning( - "Failed to enqueue Agent backend cleanup for conversation deletion: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", - stored_session.scope.tenant_id, - stored_session.scope.app_id, - stored_session.scope.conversation_id, - stored_session.scope.agent_id, - stored_session.backend_run_id, - exc_info=True, - ) - finally: - try: - session_store.mark_cleaned( - scope=stored_session.scope, - backend_run_id=stored_session.backend_run_id, - ) - except Exception: - logger.warning( - "Failed to retire Agent App runtime session for conversation deletion: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", - stored_session.scope.tenant_id, - stored_session.scope.app_id, - stored_session.scope.conversation_id, - stored_session.scope.agent_id, - stored_session.backend_run_id, - exc_info=True, - ) - + if binding_id is not None: + retired_binding_id = AgentWorkspaceService.retire_binding( + session=session, + tenant_id=app_model.tenant_id, + binding_id=binding_id, + ) + if retired_binding_id is None: + raise AgentWorkspaceNotFoundError("Conversation participant Binding is unavailable") session.delete(conversation) session.commit() - - delete_conversation_related_data.delay(conversation.id) - except Exception: session.rollback() raise + if retired_binding_id is not None: + enqueue_agent_resource_collection( + tenant_id=app_model.tenant_id, + binding_ids=(retired_binding_id,), + ) + delete_conversation_related_data.delay(conversation.id) @classmethod def get_conversational_variable( diff --git a/api/services/data_migration/import_service.py b/api/services/data_migration/import_service.py index 6b999119b2e..88728faa84f 100644 --- a/api/services/data_migration/import_service.py +++ b/api/services/data_migration/import_service.py @@ -26,6 +26,7 @@ from models import Account, ApiToken, Tenant, TenantAccountJoin, TenantAccountRo from models.enums import ApiTokenType from models.model import App from models.tools import ApiToolProvider, MCPToolProvider, WorkflowToolProvider +from services.agent.retirement_service import WorkflowAgentRetirementService from services.app_dsl_service import AppDslService from services.data_migration.dependency_discovery_service import DependencyDiscoveryService from services.data_migration.entities import ( @@ -47,6 +48,7 @@ from services.tools.api_tools_manage_service import ApiToolManageService from services.tools.mcp_tools_manage_service import MCPToolManageService from services.tools.workflow_tools_manage_service import WorkflowToolManageService from services.workflow_service import WorkflowService +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection @dataclass(frozen=True) @@ -711,7 +713,7 @@ class MigrationImportService: raise MigrationDataError(f"Referenced workflow app was not found in target tenant: {app_id}") if account_in_session is None: raise MigrationDataError(f"Operator account not found: {account.id}") - workflow = workflow_service.publish_workflow( + workflow, retirement_candidates = workflow_service.publish_workflow( session=session, app_model=app_in_session, account=account_in_session, @@ -721,6 +723,16 @@ class MigrationImportService: app_in_session.workflow_id = workflow.id app_in_session.updated_by = account.id app_in_session.updated_at = naive_utc_now() + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=target.tenant_id, + agent_ids=retirement_candidates, + account_id=account.id, + ) + enqueue_agent_resource_collection( + tenant_id=target.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) def _import_mcp_tools( self, diff --git a/api/services/snippet_dsl_service.py b/api/services/snippet_dsl_service.py index 14ac9a21f6b..0c694151252 100644 --- a/api/services/snippet_dsl_service.py +++ b/api/services/snippet_dsl_service.py @@ -19,12 +19,14 @@ from models import Account from models.snippet import CustomizedSnippet, SnippetType from models.workflow import Workflow from services.agent.dsl_service import AgentDslService +from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.workflow_publish_service import WorkflowAgentPublishService from services.dsl_content import DSL_MAX_SIZE, dsl_content_size from services.dsl_version import check_version_compatibility from services.entities.dsl_entities import CheckDependenciesResult, DslImportWarning, ImportMode, ImportStatus from services.plugin.dependencies_analysis import DependenciesAnalysisService from services.snippet_service import SNIPPET_FORBIDDEN_NODE_TYPES, SnippetService +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -426,6 +428,7 @@ class SnippetDslService: self._session.flush() # Create or update draft workflow + retirement_candidates: set[str] = set() if workflow_data: graph = workflow_data.get("graph", {}) raw_agent_packages = data.get("agent_packages") or {} @@ -444,10 +447,10 @@ class SnippetDslService: unique_hash=unique_hash, account=account, input_fields=input_fields, - sync_agent_bindings=not raw_agent_packages, + sync_agent_bindings=False, ) if raw_agent_packages: - _, warnings = AgentDslService(self._session).import_workflow_packages( + _, warnings, retirement_candidates = AgentDslService(self._session).import_workflow_packages( workflow=draft_workflow, portable_graph=graph, raw_packages=raw_agent_packages, @@ -458,8 +461,29 @@ class SnippetDslService: session=self._session, draft_workflow=draft_workflow, ) + else: + retirement_candidates = WorkflowAgentPublishService.sync_agent_bindings_for_draft( + session=self._session, + draft_workflow=draft_workflow, + account_id=account.id, + ) + WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync( + session=self._session, + draft_workflow=draft_workflow, + ) self._session.commit() + if workflow_data: + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=snippet.tenant_id, + agent_ids=retirement_candidates, + account_id=account.id, + ) + enqueue_agent_resource_collection( + tenant_id=snippet.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) return snippet def export_snippet_dsl(self, snippet: CustomizedSnippet, include_secret: bool = False) -> str: diff --git a/api/services/snippet_service.py b/api/services/snippet_service.py index c23895a4d60..1cd64216211 100644 --- a/api/services/snippet_service.py +++ b/api/services/snippet_service.py @@ -35,6 +35,7 @@ from models.workflow import ( WorkflowType, ) from repositories.factory import DifyAPIRepositoryFactory +from services.agent.retirement_service import WorkflowAgentRetirementService from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.tag_service import TagService from services.workflow_node_execution_trace_service import ( @@ -42,6 +43,7 @@ from services.workflow_node_execution_trace_service import ( assemble_workflow_node_execution_traces, ) from services.workflow_restore import apply_published_workflow_snapshot_to_draft +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection logger = logging.getLogger(__name__) @@ -80,7 +82,7 @@ class SnippetService: @contextmanager def _session_scope(self) -> Generator[Session, None, None]: - current_session = getattr(self, "_session", None) + current_session = self._session if current_session is not None: yield current_session return @@ -89,7 +91,7 @@ class SnippetService: yield session def _commit_if_owned(self, session: Session) -> None: - if getattr(self, "_session", None) is None: + if self._session is None: session.commit() @staticmethod @@ -600,12 +602,13 @@ class SnippetService: from services.agent.workflow_publish_service import WorkflowAgentPublishService + retirement_candidates: set[str] = set() with self._session_scope() as session: session.add(workflow) session.add(snippet) if sync_agent_bindings: session.flush() - WorkflowAgentPublishService.sync_agent_bindings_for_draft( + retirement_candidates = WorkflowAgentPublishService.sync_agent_bindings_for_draft( session=session, draft_workflow=workflow, account_id=account.id, @@ -615,6 +618,17 @@ class SnippetService: draft_workflow=workflow, ) self._commit_if_owned(session) + if self._session is None: + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=snippet.tenant_id, + agent_ids=retirement_candidates, + account_id=account.id, + ) + enqueue_agent_resource_collection( + tenant_id=snippet.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) return workflow def restore_published_workflow_to_draft( @@ -656,13 +670,24 @@ class SnippetService: session.flush() from services.agent.workflow_publish_service import WorkflowAgentPublishService - WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( + retirement_candidates = WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( session=session, source_workflow=source_workflow, draft_workflow=draft_workflow, account_id=account.id, ) self._commit_if_owned(session) + if self._session is None: + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=snippet.tenant_id, + agent_ids=retirement_candidates, + account_id=account.id, + ) + enqueue_agent_resource_collection( + tenant_id=snippet.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) return draft_workflow def publish_workflow( @@ -671,7 +696,7 @@ class SnippetService: session: Session, snippet: CustomizedSnippet, account: Account, - ) -> Workflow: + ) -> tuple[Workflow, set[str]]: """ Publish the draft workflow as a new version. @@ -715,7 +740,7 @@ class SnippetService: kind=WorkflowKind.SNIPPET.value, ) session.add(workflow) - WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + retirement_candidates = WorkflowAgentPublishService.copy_agent_node_bindings_to_published( session=session, draft_workflow=draft_workflow, published_workflow=workflow, @@ -728,7 +753,7 @@ class SnippetService: snippet.updated_by = account.id session.add(snippet) - return workflow + return workflow, retirement_candidates def get_all_published_workflows( self, diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index 0d73acfbca2..b1f29c2ef7e 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -76,6 +76,7 @@ from models.model import App, AppMode from models.tools import WorkflowToolProvider from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowType from repositories.factory import DifyAPIRepositoryFactory +from services.agent.retirement_service import WorkflowAgentRetirementService from services.billing_service import BillingService from services.errors.app import ( IsDraftWorkflowError, @@ -83,6 +84,7 @@ from services.errors.app import ( WorkflowHashNotEqualError, WorkflowNotFoundError, ) +from tasks.collect_agent_resources_task import enqueue_agent_resource_collection @dataclass(frozen=True) @@ -374,8 +376,9 @@ class WorkflowService: from services.agent.workflow_publish_service import WorkflowAgentPublishService session.flush() + retirement_candidates: set[str] = set() if sync_agent_bindings: - WorkflowAgentPublishService.sync_agent_bindings_for_draft( + retirement_candidates = WorkflowAgentPublishService.sync_agent_bindings_for_draft( session=session, draft_workflow=workflow, account_id=account.id, @@ -388,6 +391,16 @@ class WorkflowService: # commit db session changes if commit: session.commit() + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=app_model.tenant_id, + agent_ids=retirement_candidates, + account_id=account.id, + ) + enqueue_agent_resource_collection( + tenant_id=app_model.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) # trigger app workflow events if commit: @@ -509,7 +522,7 @@ class WorkflowService: from services.agent.workflow_publish_service import WorkflowAgentPublishService session.flush() - WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( + retirement_candidates = WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( session=session, source_workflow=source_workflow, draft_workflow=draft_workflow, @@ -517,6 +530,16 @@ class WorkflowService: ) session.commit() + binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=app_model.tenant_id, + agent_ids=retirement_candidates, + account_id=account.id, + ) + enqueue_agent_resource_collection( + tenant_id=app_model.tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + ) app_draft_workflow_was_synced.send(app_model, synced_draft_workflow=draft_workflow) return draft_workflow @@ -529,7 +552,7 @@ class WorkflowService: account: Account, marked_name: str = "", marked_comment: str = "", - ) -> Workflow: + ) -> tuple[Workflow, set[str]]: draft_workflow_stmt = select(Workflow).where( Workflow.tenant_id == app_model.tenant_id, Workflow.app_id == app_model.id, @@ -588,7 +611,7 @@ class WorkflowService: # commit db session changes session.add(workflow) - WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + retirement_candidates = WorkflowAgentPublishService.copy_agent_node_bindings_to_published( session=session, draft_workflow=draft_workflow, published_workflow=workflow, @@ -598,7 +621,7 @@ class WorkflowService: app_published_workflow_was_updated.send(app_model, published_workflow=workflow) # return new workflow - return workflow + return workflow, retirement_candidates def _validate_workflow_credentials(self, workflow: Workflow, *, session: Session) -> None: """ diff --git a/api/tasks/agent_backend_session_cleanup_task.py b/api/tasks/agent_backend_session_cleanup_task.py deleted file mode 100644 index 0dbc32b45e6..00000000000 --- a/api/tasks/agent_backend_session_cleanup_task.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Celery tasks that execute Agent backend lifecycle-only session cleanup.""" - -from __future__ import annotations - -import logging - -from celery import shared_task - -from clients.agent_backend.factory import create_agent_backend_run_client -from clients.agent_backend.request_builder import AgentBackendRunRequestBuilder -from clients.agent_backend.session_cleanup import ( - AgentBackendSessionCleanupPayload, - cleanup_agent_backend_session, -) -from configs import dify_config - -logger = logging.getLogger(__name__) - - -def _create_agent_backend_client(): - if not (dify_config.AGENT_BACKEND_USE_FAKE or dify_config.AGENT_BACKEND_BASE_URL): - return None - return create_agent_backend_run_client( - base_url=dify_config.AGENT_BACKEND_BASE_URL, - api_token=dify_config.AGENT_BACKEND_API_TOKEN, - use_fake=dify_config.AGENT_BACKEND_USE_FAKE, - fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO, - stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS, - stream_max_reconnects=dify_config.AGENT_BACKEND_STREAM_MAX_RECONNECTS, - stream_run_timeout_seconds=dify_config.AGENT_BACKEND_RUN_TIMEOUT_SECONDS, - ) - - -def _run_cleanup_task(payload_dict: dict[str, object]) -> None: - payload = AgentBackendSessionCleanupPayload.model_validate(payload_dict) - result = cleanup_agent_backend_session( - payload=payload, - client=_create_agent_backend_client(), - request_builder=AgentBackendRunRequestBuilder(), - ) - if result.status == "succeeded": - return - - log_fields = { - "tenant_id": payload.metadata.get("tenant_id"), - "app_id": payload.metadata.get("app_id"), - "workflow_run_id": payload.metadata.get("workflow_run_id"), - "node_id": payload.metadata.get("node_id"), - "conversation_id": payload.metadata.get("conversation_id"), - "agent_id": payload.metadata.get("agent_id"), - "previous_agent_backend_run_id": payload.metadata.get("previous_agent_backend_run_id"), - "failed_agent_backend_run_id": payload.metadata.get("failed_agent_backend_run_id"), - "cleanup_run_id": result.cleanup_run_id, - "reason": result.reason, - } - if result.status == "skipped": - logger.info("Agent backend session cleanup skipped: %s", log_fields) - return - - logger.warning("Agent backend session cleanup failed: %s", log_fields) - - -@shared_task(queue="workflow_storage") -def cleanup_workflow_agent_runtime_session(payload_dict: dict[str, object]) -> None: - """Run one workflow-owned Agent backend cleanup payload.""" - _run_cleanup_task(payload_dict) - - -@shared_task(queue="conversation") -def cleanup_conversation_agent_runtime_session(payload_dict: dict[str, object]) -> None: - """Run one conversation-owned Agent backend cleanup payload.""" - _run_cleanup_task(payload_dict) diff --git a/api/tasks/app_generate/resume_agent_app_task.py b/api/tasks/app_generate/resume_agent_app_task.py index 0b884b014bc..eac6d273425 100644 --- a/api/tasks/app_generate/resume_agent_app_task.py +++ b/api/tasks/app_generate/resume_agent_app_task.py @@ -51,6 +51,7 @@ def resume_agent_app_execution(*, conversation_id: str, form_id: str) -> None: app_model=app_model, user=user, conversation_id=conversation_id, + form_id=form_id, invoke_from=_resolve_invoke_from(conversation), session=db.session(), ) diff --git a/api/tasks/collect_agent_resources_task.py b/api/tasks/collect_agent_resources_task.py new file mode 100644 index 00000000000..2b84cc01b79 --- /dev/null +++ b/api/tasks/collect_agent_resources_task.py @@ -0,0 +1,75 @@ +"""Asynchronously collect retired Agent working resources.""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable + +from celery import shared_task + +from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.workspace_service import AgentWorkspaceService + +logger = logging.getLogger(__name__) + + +@shared_task(queue="retention") +def collect_agent_resources( + *, + tenant_id: str, + binding_ids: list[str], + workspace_ids: list[str], + home_snapshot_ids: list[str], +) -> None: + """Collect only the explicitly identified RETIRED resources.""" + + collectors = ( + (workspace_ids, "workspace_id", AgentWorkspaceService.collect_retired_workspace), + (binding_ids, "binding_id", AgentWorkspaceService.collect_retired_binding), + ( + home_snapshot_ids, + "home_snapshot_id", + AgentHomeSnapshotService.collect_retired_home_snapshot, + ), + ) + for resource_ids, argument_name, collector in collectors: + for resource_id in resource_ids: + try: + collector(tenant_id=tenant_id, **{argument_name: resource_id}) + except Exception: + logger.exception( + "Failed to collect retired Agent resource", + extra={ + "tenant_id": tenant_id, + "resource_type": argument_name.removesuffix("_id"), + "resource_id": resource_id, + }, + ) + + +def enqueue_agent_resource_collection( + *, + tenant_id: str, + binding_ids: Iterable[str] = (), + workspace_ids: Iterable[str] = (), + home_snapshot_ids: Iterable[str] = (), +) -> None: + """Best-effort enqueue of physical collection after retire has committed.""" + + payload = { + "binding_ids": sorted({resource_id for resource_id in binding_ids if resource_id}), + "workspace_ids": sorted({resource_id for resource_id in workspace_ids if resource_id}), + "home_snapshot_ids": sorted({resource_id for resource_id in home_snapshot_ids if resource_id}), + } + if not any(payload.values()): + return + try: + collect_agent_resources.delay(tenant_id=tenant_id, **payload) + except Exception: + logger.exception( + "Failed to enqueue retired Agent resource collection", + extra={"tenant_id": tenant_id, **payload}, + ) + + +__all__ = ["collect_agent_resources", "enqueue_agent_resource_collection"] diff --git a/api/tasks/remove_app_and_related_data_task.py b/api/tasks/remove_app_and_related_data_task.py index 2010deb4a80..a4fb6d57207 100644 --- a/api/tasks/remove_app_and_related_data_task.py +++ b/api/tasks/remove_app_and_related_data_task.py @@ -5,25 +5,17 @@ from typing import Any, cast import click import sqlalchemy as sa -from agenton.compositor import CompositorSessionSnapshot from celery import shared_task -from dify_agent.protocol import RuntimeLayerSpec -from pydantic import JsonValue, TypeAdapter from sqlalchemy import delete, select from sqlalchemy.engine import CursorResult from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import sessionmaker -from clients.agent_backend.session_cleanup import AgentBackendSessionCleanupPayload from configs import dify_config from core.db.session_factory import session_factory from extensions.ext_database import db from libs.archive_storage import ArchiveStorageNotConfiguredError, get_archive_storage -from libs.datetime_utils import naive_utc_now from models import ( - AgentRuntimeSession, - AgentRuntimeSessionOwnerType, - AgentRuntimeSessionStatus, ApiToken, AppAnnotationHitHistory, AppAnnotationSetting, @@ -58,13 +50,8 @@ from models.workflow import ( ) from repositories.factory import DifyAPIRepositoryFactory from services.api_token_service import ApiTokenCache -from tasks.agent_backend_session_cleanup_task import ( - cleanup_conversation_agent_runtime_session, - cleanup_workflow_agent_runtime_session, -) logger = logging.getLogger(__name__) -_RUNTIME_LAYER_SPECS_ADAPTER: TypeAdapter[list[RuntimeLayerSpec]] = TypeAdapter(list[RuntimeLayerSpec]) @shared_task(queue="app_deletion", bind=True, max_retries=3) @@ -72,7 +59,6 @@ def remove_app_and_related_data_task(self, tenant_id: str, app_id: str): logger.info(click.style(f"Start deleting app and related data: {tenant_id}:{app_id}", fg="green")) start_at = time.perf_counter() try: - _cleanup_active_agent_runtime_sessions_for_app(tenant_id, app_id) # Delete related data _delete_app_model_configs(tenant_id, app_id) _delete_app_site(tenant_id, app_id) @@ -113,143 +99,6 @@ def remove_app_and_related_data_task(self, tenant_id: str, app_id: str): raise self.retry(exc=e, countdown=60) # Retry after 60 seconds -def _cleanup_active_agent_runtime_sessions_for_app(tenant_id: str, app_id: str, *, batch_size: int = 100) -> None: - """Best-effort fan-out for ACTIVE Agent runtime sessions during app deletion. - - App deletion must not block on synchronous Agent backend lifecycle work, so - this helper scans ACTIVE ``agent_runtime_sessions`` rows in batches, - dispatches owner-specific cleanup tasks only when enough persisted data - exists to replay a lifecycle-only run, and then marks each visited row - ``CLEANED`` locally regardless of enqueue outcome. The local retirement is - the contract that lets the rest of app deletion continue even when backend - cleanup dispatch is skipped or fails. - """ - if batch_size <= 0: - raise ValueError("batch_size must be positive") - - while True: - with session_factory.create_session() as session: - row_ids = session.scalars( - select(AgentRuntimeSession.id) - .where( - AgentRuntimeSession.tenant_id == tenant_id, - AgentRuntimeSession.app_id == app_id, - AgentRuntimeSession.status == AgentRuntimeSessionStatus.ACTIVE, - ) - .order_by(AgentRuntimeSession.updated_at.asc()) - .limit(batch_size) - ).all() - - if not row_ids: - return - - retired_count = 0 - for row_id in row_ids: - with session_factory.create_session() as session: - row = session.get(AgentRuntimeSession, row_id) - if row is None or row.status != AgentRuntimeSessionStatus.ACTIVE: - retired_count += 1 - continue - - try: - payload = _build_agent_runtime_session_cleanup_payload(row) - if payload is not None: - _enqueue_agent_runtime_session_cleanup(row=row, payload=payload) - except Exception: - logger.warning( - "Failed to enqueue Agent backend cleanup during app deletion: " - "tenant_id=%s app_id=%s owner_type=%s conversation_id=%s workflow_run_id=%s " - "node_id=%s agent_id=%s backend_run_id=%s", - row.tenant_id, - row.app_id, - row.owner_type, - row.conversation_id, - row.workflow_run_id, - row.node_id, - row.agent_id, - row.backend_run_id, - exc_info=True, - ) - finally: - try: - row.status = AgentRuntimeSessionStatus.CLEANED - row.cleaned_at = naive_utc_now() - session.commit() - retired_count += 1 - except Exception: - session.rollback() - logger.warning( - "Failed to retire Agent runtime session during app deletion: " - "tenant_id=%s app_id=%s owner_type=%s conversation_id=%s workflow_run_id=%s " - "node_id=%s agent_id=%s backend_run_id=%s", - row.tenant_id, - row.app_id, - row.owner_type, - row.conversation_id, - row.workflow_run_id, - row.node_id, - row.agent_id, - row.backend_run_id, - exc_info=True, - ) - - if retired_count == 0: - logger.warning( - "Failed to retire any active Agent runtime sessions during app deletion: tenant_id=%s app_id=%s", - tenant_id, - app_id, - ) - return - - -def _build_agent_runtime_session_cleanup_payload( - row: AgentRuntimeSession, -) -> AgentBackendSessionCleanupPayload | None: - runtime_layer_specs = _RUNTIME_LAYER_SPECS_ADAPTER.validate_json(row.composition_layer_specs or "[]") - if not runtime_layer_specs: - return None - - metadata: dict[str, JsonValue] = { - "tenant_id": row.tenant_id, - "app_id": row.app_id, - "agent_id": row.agent_id, - "agent_config_snapshot_id": row.agent_config_snapshot_id, - "previous_agent_backend_run_id": row.backend_run_id, - } - if row.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION: - metadata["conversation_id"] = row.conversation_id - idempotency_key = ( - f"{row.tenant_id}:{row.app_id}:{row.conversation_id}:" - f"{row.agent_id}:app-delete-cleanup:{row.id or row.backend_run_id or 'no-session-id'}" - ) - else: - metadata["workflow_run_id"] = row.workflow_run_id - metadata["node_id"] = row.node_id - idempotency_key = ( - f"{row.tenant_id}:{row.app_id}:{row.workflow_run_id}:{row.node_id}:" - f"{row.agent_id}:app-delete-cleanup:{row.id or row.backend_run_id or 'no-session-id'}" - ) - - return AgentBackendSessionCleanupPayload( - session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), - runtime_layer_specs=runtime_layer_specs, - idempotency_key=idempotency_key, - metadata=metadata, - ) - - -def _enqueue_agent_runtime_session_cleanup( - *, - row: AgentRuntimeSession, - payload: AgentBackendSessionCleanupPayload, -) -> None: - payload_dict = payload.model_dump(mode="json") - if row.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION: - cleanup_conversation_agent_runtime_session.delay(payload_dict) - return - cleanup_workflow_agent_runtime_session.delay(payload_dict) - - def _delete_app_model_configs(tenant_id: str, app_id: str): def del_model_config(session, model_config_id: str): session.execute( diff --git a/api/tasks/workflow_node_execution_tasks.py b/api/tasks/workflow_node_execution_tasks.py index 0d5475a56d2..e55dcf45330 100644 --- a/api/tasks/workflow_node_execution_tasks.py +++ b/api/tasks/workflow_node_execution_tasks.py @@ -13,6 +13,7 @@ from celery import shared_task from sqlalchemy import select from core.db.session_factory import session_factory +from core.workflow.node_execution_process_data import preserve_workflow_agent_binding_id from graphon.entities.workflow_node_execution import ( WorkflowNodeExecution, ) @@ -144,8 +145,9 @@ def _update_node_execution_from_domain(node_execution: WorkflowNodeExecutionMode # Update serialized data json_converter = WorkflowRuntimeTypeConverter() node_execution.inputs = json.dumps(json_converter.to_json_encodable(execution.inputs)) if execution.inputs else "{}" + process_data = preserve_workflow_agent_binding_id(node_execution.process_data_dict, execution.process_data) node_execution.process_data = ( - json.dumps(json_converter.to_json_encodable(execution.process_data)) if execution.process_data else "{}" + json.dumps(json_converter.to_json_encodable(process_data)) if process_data is not None else "{}" ) node_execution.outputs = ( json.dumps(json_converter.to_json_encodable(execution.outputs)) if execution.outputs else "{}" diff --git a/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py b/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py index 32591c13312..bf84d206da2 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py +++ b/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py @@ -124,6 +124,7 @@ class TestAppDslService: patch("services.app_service.ModelManager.for_tenant") as mock_model_manager, patch("services.app_service.FeatureService") as mock_feature_service, patch("services.app_service.EnterpriseService") as mock_enterprise_service, + patch("services.agent.home_snapshot_service.AgentHomeSnapshotService._client") as mock_home_snapshot_client, ): mock_workflow_service.return_value.get_draft_workflow.return_value = None mock_workflow_service.return_value.sync_draft_workflow.return_value = MagicMock() @@ -142,6 +143,9 @@ class TestAppDslService: mock_feature_service.get_system_features.return_value.webapp_auth.enabled = False mock_enterprise_service.WebAppAuth.update_app_access_mode.return_value = None mock_enterprise_service.WebAppAuth.cleanup_webapp.return_value = None + mock_home_snapshot_client.return_value.__enter__.return_value.initialize_home_snapshot_sync.side_effect = ( + lambda request: SimpleNamespace(snapshot_ref=f"test:{request.home_snapshot_id}") + ) yield { "workflow_service": mock_workflow_service, @@ -1034,7 +1038,9 @@ class TestAppDslService: ) ) - imported_graph, warnings = AgentDslService(db_session_with_containers).import_workflow_packages( + imported_graph, warnings, retirement_candidates = AgentDslService( + db_session_with_containers + ).import_workflow_packages( workflow=workflow, portable_graph=graph, raw_packages={"agent_1": package.model_dump(mode="json")}, @@ -1043,6 +1049,7 @@ class TestAppDslService: db_session_with_containers.commit() assert warnings == [] + assert retirement_candidates == set() graph_bindings = [node["data"]["agent_binding"] for node in imported_graph["nodes"]] assert all(binding["binding_type"] == WorkflowAgentBindingType.INLINE_AGENT.value for binding in graph_bindings) assert len({binding["agent_id"] for binding in graph_bindings}) == 2 diff --git a/api/tests/test_containers_integration_tests/services/test_app_service.py b/api/tests/test_containers_integration_tests/services/test_app_service.py index 2875f014380..599f7565db0 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_service.py +++ b/api/tests/test_containers_integration_tests/services/test_app_service.py @@ -1,4 +1,5 @@ from datetime import datetime +from types import SimpleNamespace from unittest.mock import create_autospec, patch import pytest @@ -29,6 +30,7 @@ class TestAppService: patch("services.app_service.EnterpriseService") as mock_enterprise_service, patch("services.app_service.ModelManager.for_tenant") as mock_model_manager, patch("services.account_service.FeatureService") as mock_account_feature_service, + patch("services.agent.home_snapshot_service.AgentHomeSnapshotService._client") as mock_home_snapshot_client, ): # Setup default mock returns for app service mock_feature_service.get_system_features.return_value.webapp_auth.enabled = False @@ -42,6 +44,9 @@ class TestAppService: mock_model_instance = mock_model_manager.return_value mock_model_instance.get_default_model_instance.return_value = None mock_model_instance.get_default_provider_model_name.return_value = ("openai", "gpt-3.5-turbo") + mock_home_snapshot_client.return_value.__enter__.return_value.initialize_home_snapshot_sync.side_effect = ( + lambda request: SimpleNamespace(snapshot_ref=f"test:{request.home_snapshot_id}") + ) yield { "feature_service": mock_feature_service, diff --git a/api/tests/test_containers_integration_tests/services/test_conversation_service.py b/api/tests/test_containers_integration_tests/services/test_conversation_service.py index 09411f0b0e1..75f017fbc50 100644 --- a/api/tests/test_containers_integration_tests/services/test_conversation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_conversation_service.py @@ -6,13 +6,11 @@ from unittest.mock import patch from uuid import uuid4 import pytest -from agenton.compositor import CompositorSessionSnapshot from sqlalchemy import select from sqlalchemy.orm import Session -from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore from core.app.entities.app_invoke_entities import InvokeFrom -from models import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus, TenantAccountRole +from models import TenantAccountRole from models.account import Account, Tenant, TenantAccountJoin from models.enums import ConversationFromSource, EndUserType from models.model import App, Conversation, EndUser, Message, MessageAnnotation @@ -1077,9 +1075,8 @@ class TestConversationServiceExport: # Assert assert result == conversation - @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") @patch("services.conversation_service.delete_conversation_related_data") - def test_delete_conversation(self, mock_delete_task, mock_cleanup_task, db_session_with_containers: Session): + def test_delete_conversation(self, mock_delete_task, db_session_with_containers: Session): """ Test conversation deletion with async cleanup. @@ -1098,20 +1095,6 @@ class TestConversationServiceExport: user, ) conversation_id = conversation.id - runtime_session = AgentRuntimeSession( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id=str(uuid4()), - agent_config_snapshot_id=str(uuid4()), - backend_run_id="backend-run-1", - session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), - composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', - conversation_id=conversation.id, - status=AgentRuntimeSessionStatus.ACTIVE, - ) - db_session_with_containers.add(runtime_session) - db_session_with_containers.commit() # Act - Delete the conversation ConversationService.delete( @@ -1126,27 +1109,11 @@ class TestConversationServiceExport: # Step 2: Async cleanup task triggered # The Celery task will handle cleanup of messages, annotations, etc. mock_delete_task.delay.assert_called_once_with(conversation_id) - mock_cleanup_task.delay.assert_called_once() - cleanup_payload = mock_cleanup_task.delay.call_args.args[0] - assert cleanup_payload["metadata"]["conversation_id"] == conversation_id - assert ( - cleanup_payload["idempotency_key"] - == f"{app_model.tenant_id}:{app_model.id}:{conversation_id}:agent-runtime-session-cleanup:" - f"{runtime_session.agent_id}:{runtime_session.agent_config_snapshot_id}:{runtime_session.backend_run_id}" - ) - runtime_session_row = db_session_with_containers.scalar( - select(AgentRuntimeSession).where(AgentRuntimeSession.id == runtime_session.id) - ) - assert runtime_session_row is not None - assert runtime_session_row.status == AgentRuntimeSessionStatus.CLEANED - - @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") @patch("services.conversation_service.delete_conversation_related_data") def test_delete_conversation_not_owned_by_account( self, mock_delete_task, - mock_cleanup_task, db_session_with_containers: Session, ): """ @@ -1178,22 +1145,17 @@ class TestConversationServiceExport: not_deleted = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation.id)) assert not_deleted is not None mock_delete_task.delay.assert_not_called() - mock_cleanup_task.delay.assert_not_called() - @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") @patch("services.conversation_service.delete_conversation_related_data") def test_delete_handles_exception_and_rollback( self, mock_delete_task, - mock_cleanup_task, db_session_with_containers: Session, ): """ Test that delete propagates exceptions and does not trigger the cleanup task. - When a DB error occurs during deletion, the conversation row stays in - place, but any already-enqueued Agent backend cleanup remains a - best-effort terminal lifecycle action. + When a DB error occurs during deletion, the conversation row stays in place. """ # Arrange app_model, user = ConversationServiceIntegrationTestDataFactory.create_app_and_account( @@ -1203,20 +1165,6 @@ class TestConversationServiceExport: db_session_with_containers, app_model, user ) conversation_id = conversation.id - runtime_session = AgentRuntimeSession( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id=str(uuid4()), - agent_config_snapshot_id=str(uuid4()), - backend_run_id="backend-run-rollback", - session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), - composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', - conversation_id=conversation.id, - status=AgentRuntimeSessionStatus.ACTIVE, - ) - db_session_with_containers.add(runtime_session) - db_session_with_containers.commit() # Act — force an error during the delete to exercise the rollback path with patch.object(db_session_with_containers, "delete", side_effect=Exception("DB error")): @@ -1228,111 +1176,9 @@ class TestConversationServiceExport: session=db_session_with_containers, ) - # Assert — related-data deletion is not scheduled, but the backend - # cleanup task was already enqueued before the row delete failed. + # Assert — related-data deletion is not scheduled. mock_delete_task.delay.assert_not_called() - mock_cleanup_task.delay.assert_called_once() - cleanup_payload = mock_cleanup_task.delay.call_args.args[0] - assert ( - cleanup_payload["idempotency_key"] - == f"{app_model.tenant_id}:{app_model.id}:{conversation_id}:agent-runtime-session-cleanup:" - f"{runtime_session.agent_id}:{runtime_session.agent_config_snapshot_id}:{runtime_session.backend_run_id}" - ) # Conversation is still present because the deletion was never committed still_there = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation_id)) assert still_there is not None - - @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") - @patch("services.conversation_service.delete_conversation_related_data") - def test_delete_ignores_mark_cleaned_failure( - self, - mock_delete_task, - mock_cleanup_task, - db_session_with_containers: Session, - ): - app_model, user = ConversationServiceIntegrationTestDataFactory.create_app_and_account( - db_session_with_containers - ) - conversation = ConversationServiceIntegrationTestDataFactory.create_conversation( - db_session_with_containers, - app_model, - user, - ) - runtime_session = AgentRuntimeSession( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id=str(uuid4()), - agent_config_snapshot_id=str(uuid4()), - backend_run_id="backend-run-cleanup-failure", - session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), - composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', - conversation_id=conversation.id, - status=AgentRuntimeSessionStatus.ACTIVE, - ) - db_session_with_containers.add(runtime_session) - db_session_with_containers.commit() - - with patch.object(AgentAppRuntimeSessionStore, "mark_cleaned", side_effect=RuntimeError("cleanup failed")): - ConversationService.delete( - app_model=app_model, - conversation_id=conversation.id, - user=user, - session=db_session_with_containers, - ) - - deleted = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation.id)) - assert deleted is None - mock_delete_task.delay.assert_called_once_with(conversation.id) - mock_cleanup_task.delay.assert_called_once() - - @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") - @patch("services.conversation_service.delete_conversation_related_data") - def test_delete_ignores_cleanup_enqueue_failure_and_still_retires_runtime_session( - self, - mock_delete_task, - mock_cleanup_task, - db_session_with_containers: Session, - ): - app_model, user = ConversationServiceIntegrationTestDataFactory.create_app_and_account( - db_session_with_containers - ) - conversation = ConversationServiceIntegrationTestDataFactory.create_conversation( - db_session_with_containers, - app_model, - user, - ) - conversation_id = conversation.id - runtime_session = AgentRuntimeSession( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id=str(uuid4()), - agent_config_snapshot_id=str(uuid4()), - backend_run_id="backend-run-enqueue-failure", - session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), - composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', - conversation_id=conversation.id, - status=AgentRuntimeSessionStatus.ACTIVE, - ) - db_session_with_containers.add(runtime_session) - db_session_with_containers.commit() - mock_cleanup_task.delay.side_effect = RuntimeError("queue down") - - ConversationService.delete( - app_model=app_model, - conversation_id=conversation_id, - user=user, - session=db_session_with_containers, - ) - - deleted = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation_id)) - assert deleted is None - mock_delete_task.delay.assert_called_once_with(conversation_id) - mock_cleanup_task.delay.assert_called_once() - runtime_session_row = db_session_with_containers.scalar( - select(AgentRuntimeSession).where(AgentRuntimeSession.id == runtime_session.id) - ) - assert runtime_session_row is not None - assert runtime_session_row.status == AgentRuntimeSessionStatus.CLEANED diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_service.py b/api/tests/test_containers_integration_tests/services/test_workflow_service.py index 6531ed4fbb0..dd188693360 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workflow_service.py @@ -870,12 +870,13 @@ class TestWorkflowService: from unittest.mock import patch with patch("flask_login.utils._get_user", return_value=account, autospec=True): - result = workflow_service.publish_workflow( + result, retirement_candidates = workflow_service.publish_workflow( session=db_session_with_containers, app_model=app, account=account ) # Assert assert result is not None + assert retirement_candidates == set() assert result.version != Workflow.VERSION_DRAFT # Version should be a timestamp format like '2025-08-22 00:10:24.722051' assert isinstance(result.version, str) diff --git a/api/tests/unit_tests/clients/agent_backend/test_cleanup_composition_compositor_integration.py b/api/tests/unit_tests/clients/agent_backend/test_cleanup_composition_compositor_integration.py deleted file mode 100644 index 00e74b1b659..00000000000 --- a/api/tests/unit_tests/clients/agent_backend/test_cleanup_composition_compositor_integration.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Integration test for the cleanup request against the real agenton compositor. - -The bug fixed by A+D was invisible to unit tests that use ``FakeAgentBackendRunClient`` -because the fake client never runs agenton's ``_validate_session_snapshot``. This -test plugs a cleanup request through the real ``Compositor`` (with the same -providers the agent backend wires in production) so that the snapshot-vs- -composition name-order check would fail loudly if the cleanup builder ever -regressed back to the empty-composition shape. -""" - -from __future__ import annotations - -from typing import cast - -import pytest -from agenton.compositor import Compositor, CompositorSessionSnapshot, LayerProvider -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from agenton_collections.layers.plain import PLAIN_PROMPT_LAYER_TYPE_ID -from agenton_collections.layers.plain.basic import PromptLayer -from agenton_collections.layers.pydantic_ai import PYDANTIC_AI_HISTORY_LAYER_TYPE_ID, PydanticAIHistoryLayer - -from clients.agent_backend import AgentBackendRunRequestBuilder, RuntimeLayerSpec - - -def test_cleanup_request_passes_agenton_snapshot_validation(): - """The cleanup request's composition layer names must match the (filtered) - snapshot's layer names exactly — agenton's compositor enforces this and - the agent backend rejects mismatches as ``run_failed`` asynchronously, - which is the trap A/D fixed.""" - # Persisted (non-plugin) layer specs — these are what cleanup will replay. - # We exclude the dify.execution_context layer from this integration check - # because its real provider needs a plugin-daemon HTTP client; the cleanup - # validation we are exercising is the snapshot-vs-composition name check, - # which is purely structural and does not depend on which non-plugin layer - # types appear. - persisted_specs = [ - RuntimeLayerSpec( - name="workflow_node_job_prompt", - type=PLAIN_PROMPT_LAYER_TYPE_ID, - config={"prefix": "Do the cleanup."}, - ), - RuntimeLayerSpec(name="history", type=PYDANTIC_AI_HISTORY_LAYER_TYPE_ID), - ] - # Saved snapshot still carries the LLM layer entry — cleanup's - # ``_filter_snapshot_to_specs`` must drop it so names match. - full_snapshot = CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="workflow_node_job_prompt", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - LayerSessionSnapshot( - name="history", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={"messages": []}, - ), - LayerSessionSnapshot( - name="llm", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - ] - ) - - cleanup_request = AgentBackendRunRequestBuilder().build_cleanup_request( - session_snapshot=full_snapshot, - runtime_layer_specs=persisted_specs, - ) - - # Drive the real agenton compositor through ``from_config`` + ``_create_run`` - # the same way the agent backend's RunScheduler does. ``_create_run`` is the - # private path that calls ``_validate_session_snapshot``; we use it directly - # to keep the test synchronous (no async ``enter()`` lifecycle needed — - # validation is the only thing under test). - config = { - "schema_version": 1, - "layers": [ - {"name": layer.name, "type": layer.type, "deps": dict(layer.deps), "metadata": dict(layer.metadata)} - for layer in cleanup_request.composition.layers - ], - } - compositor = Compositor.from_config( - config, - providers=[ - LayerProvider.from_layer_type(PromptLayer), - LayerProvider.from_layer_type(PydanticAIHistoryLayer), - ], - ) - - layer_configs = {layer.name: layer.config for layer in cleanup_request.composition.layers} - # This is the call that would raise ``ValueError`` if the cleanup snapshot - # and composition disagreed on layer names — the exact failure mode the - # original ``layers=[]`` cleanup hit. - run = compositor._create_run( # type: ignore[reportPrivateUsage] - configs=cast(dict[str, object], layer_configs), - session_snapshot=cleanup_request.session_snapshot, - ) - assert list(run.slots.keys()) == ["workflow_node_job_prompt", "history"] - - -def test_cleanup_request_with_mismatched_specs_would_be_rejected_by_agenton(): - """Regression sentinel: if a future refactor stops filtering the snapshot, - agenton would reject the request — and that rejection is what the runtime - fix is preventing. We confirm the validator does fail when given the - pre-fix shape so the previous test's success is not a coincidence.""" - snapshot_with_extra = CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="history", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - LayerSessionSnapshot( - name="llm", # extra layer not in composition - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - ] - ) - compositor = Compositor.from_config( - { - "schema_version": 1, - "layers": [{"name": "history", "type": PYDANTIC_AI_HISTORY_LAYER_TYPE_ID, "deps": {}, "metadata": {}}], - }, - providers=[LayerProvider.from_layer_type(PydanticAIHistoryLayer)], - ) - - with pytest.raises(ValueError, match="layer names must match"): - compositor._create_run( # type: ignore[reportPrivateUsage] - configs={}, - session_snapshot=snapshot_with_extra, - ) diff --git a/api/tests/unit_tests/clients/agent_backend/test_client.py b/api/tests/unit_tests/clients/agent_backend/test_client.py index 421ffad9092..cf13067d239 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_client.py +++ b/api/tests/unit_tests/clients/agent_backend/test_client.py @@ -40,6 +40,7 @@ def _request(): agent_mode="workflow_run", invoke_from="debugger", ), + backend_binding_ref="binding-ref-1", workflow_node_job_prompt="Do the task.", user_prompt="hello", ) diff --git a/api/tests/unit_tests/clients/agent_backend/test_fake_client.py b/api/tests/unit_tests/clients/agent_backend/test_fake_client.py index 5862117f622..6ad82633036 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_fake_client.py +++ b/api/tests/unit_tests/clients/agent_backend/test_fake_client.py @@ -23,6 +23,7 @@ def _request(): agent_mode="workflow_run", invoke_from="debugger", ), + backend_binding_ref="binding-ref-1", workflow_node_job_prompt="Do the task.", user_prompt="hello", ) diff --git a/api/tests/unit_tests/clients/agent_backend/test_request_builder.py b/api/tests/unit_tests/clients/agent_backend/test_request_builder.py index b2b77275fa7..afbd5c195f2 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_request_builder.py +++ b/api/tests/unit_tests/clients/agent_backend/test_request_builder.py @@ -1,10 +1,7 @@ from typing import Any, cast import pytest -from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot from agenton.layers import ExitIntent -from agenton.layers.base import LifecycleState from agenton_collections.layers.plain import PLAIN_PROMPT_LAYER_TYPE_ID, PromptLayerConfig from agenton_collections.layers.pydantic_ai import PYDANTIC_AI_HISTORY_LAYER_TYPE_ID from dify_agent.layers.dify_core_tools import ( @@ -45,8 +42,6 @@ from clients.agent_backend import ( AgentBackendOutputConfig, AgentBackendRunRequestBuilder, AgentBackendWorkflowNodeRunInput, - RuntimeLayerSpec, - extract_runtime_layer_specs, redact_for_agent_backend_log, ) from clients.agent_backend.request_builder import DIFY_DRIVE_LAYER_ID, DIFY_SHELL_LAYER_ID @@ -71,6 +66,7 @@ def _run_input() -> AgentBackendWorkflowNodeRunInput: agent_mode="workflow_run", invoke_from="debugger", ), + backend_binding_ref="binding-ref-1", idempotency_key="workflow-run-1:node-execution-1", agent_soul_prompt="You are a careful reviewer.", workflow_node_job_prompt="Review the previous node output.", @@ -271,89 +267,6 @@ def test_request_builder_adds_knowledge_layer_when_configured(): assert knowledge_config.sets[0].dataset_ids == ["dataset-1"] -def test_request_builder_can_delete_on_exit_for_cleanup_paths(): - run_input = _run_input() - run_input.suspend_on_exit = False - - request = AgentBackendRunRequestBuilder().build_for_workflow_node(run_input) - - assert request.on_exit.default is ExitIntent.DELETE - - -def test_request_builder_builds_cleanup_request_replays_persisted_layer_specs(): - """The cleanup request must replay the persisted (non-plugin) layer specs - and filter the snapshot to match so the agenton compositor's - snapshot-vs-composition name-order validator passes.""" - session_snapshot = CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot(name="history", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={"k": 1}), - LayerSessionSnapshot(name="llm", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), - ] - ) - specs = [RuntimeLayerSpec(name="history", type="pydantic_ai.history")] - - request = AgentBackendRunRequestBuilder().build_cleanup_request( - session_snapshot=session_snapshot, - runtime_layer_specs=specs, - idempotency_key="run-1:node-1:binding-1:agent-session-cleanup", - metadata={"workflow_run_id": "run-1"}, - ) - - assert [layer.name for layer in request.composition.layers] == ["history"] - assert request.session_snapshot is not None - assert [layer.name for layer in request.session_snapshot.layers] == ["history"] - assert request.on_exit.default is ExitIntent.DELETE - assert request.idempotency_key == "run-1:node-1:binding-1:agent-session-cleanup" - assert request.metadata["agent_backend_lifecycle"] == "session_cleanup" - assert "purpose" not in request.model_dump(mode="json") - - -def test_request_builder_rejects_empty_runtime_layer_specs(): - """Empty specs would put us back in the original ``layers=[]`` trap that - fails on agenton's snapshot-vs-composition validation.""" - with pytest.raises(ValueError, match="runtime_layer_specs"): - AgentBackendRunRequestBuilder().build_cleanup_request( - session_snapshot=CompositorSessionSnapshot(layers=[]), - runtime_layer_specs=[], - ) - - -def test_extract_runtime_layer_specs_drops_plugin_layers_keeps_configs(): - from dify_agent.protocol import RunComposition, RunLayerSpec - - composition = RunComposition( - layers=[ - RunLayerSpec( - name="agent_soul_prompt", - type="plain.prompt", - config=PromptLayerConfig(prefix="hello"), - ), - RunLayerSpec( - name="llm", - type="dify.plugin.llm", - config=None, # protocol allows None; the redacted config is what matters - ), - RunLayerSpec( - name="tools", - type="dify.plugin.tools", - ), - RunLayerSpec( - name="history", - type="pydantic_ai.history", - ), - ] - ) - - specs = extract_runtime_layer_specs(composition) - - assert [spec.name for spec in specs] == ["agent_soul_prompt", "history"] - # Non-plugin configs are dumped as JSON-compatible dicts so the persisted - # row can be replayed without holding live pydantic instances. - soul_config = specs[0].config - assert isinstance(soul_config, dict) - assert soul_config.get("prefix") == "hello" - - def test_request_builder_rejects_blank_prompts(): with pytest.raises(ValidationError): AgentBackendWorkflowNodeRunInput( @@ -397,6 +310,7 @@ def _agent_app_input(*, include_shell: bool = False) -> AgentBackendAgentAppRunI agent_mode="agent_app", invoke_from="web-app", ), + backend_binding_ref="binding-ref-1", agent_soul_prompt="You are Iris.", user_prompt="List files.", include_shell=include_shell, @@ -422,7 +336,7 @@ def test_workflow_request_builder_adds_shell_layer_when_include_shell(): assert shell.type == DIFY_SHELL_LAYER_TYPE_ID # The shell layer depends on execution_context so the agent server can mint # per-command Agent Stub env for sandbox CLI forwarding. - assert shell.deps == {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID} + assert shell.deps == {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID, "runtime": "runtime"} shell_config = cast(DifyShellLayerConfig, shell.config) assert shell_config.env[0].name == "PROJECT_NAME" @@ -436,7 +350,10 @@ def test_workflow_request_builder_binds_drive_to_shell_when_configured(): layers = {layer.name: layer for layer in request.composition.layers} layer_names = [layer.name for layer in request.composition.layers] - assert layers[DIFY_SHELL_LAYER_ID].deps == {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID} + assert layers[DIFY_SHELL_LAYER_ID].deps == { + "execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID, + "runtime": "runtime", + } shell_config = cast(DifyShellLayerConfig, layers[DIFY_SHELL_LAYER_ID].config) assert shell_config.agent_stub_drive_ref == "agent-agent-1" assert layers[DIFY_DRIVE_LAYER_ID].deps == {"shell": DIFY_SHELL_LAYER_ID} @@ -470,7 +387,10 @@ def test_agent_app_request_builder_adds_shell_layer_when_include_shell(): assert DIFY_SHELL_LAYER_ID in layers assert layers[DIFY_SHELL_LAYER_ID].type == DIFY_SHELL_LAYER_TYPE_ID - assert layers[DIFY_SHELL_LAYER_ID].deps == {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID} + assert layers[DIFY_SHELL_LAYER_ID].deps == { + "execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID, + "runtime": "runtime", + } shell_config = cast(DifyShellLayerConfig, layers[DIFY_SHELL_LAYER_ID].config) assert shell_config.env[0].name == "APP_ENV" @@ -483,7 +403,10 @@ def test_agent_app_request_builder_binds_drive_to_shell_when_configured(): layers = {layer.name: layer for layer in request.composition.layers} layer_names = [layer.name for layer in request.composition.layers] - assert layers[DIFY_SHELL_LAYER_ID].deps == {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID} + assert layers[DIFY_SHELL_LAYER_ID].deps == { + "execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID, + "runtime": "runtime", + } shell_config = cast(DifyShellLayerConfig, layers[DIFY_SHELL_LAYER_ID].config) assert shell_config.agent_stub_drive_ref == "agent-agent-1" assert layers[DIFY_DRIVE_LAYER_ID].deps == {"shell": DIFY_SHELL_LAYER_ID} diff --git a/api/tests/unit_tests/clients/agent_backend/test_session_cleanup.py b/api/tests/unit_tests/clients/agent_backend/test_session_cleanup.py deleted file mode 100644 index 6b72850aeca..00000000000 --- a/api/tests/unit_tests/clients/agent_backend/test_session_cleanup.py +++ /dev/null @@ -1,124 +0,0 @@ -from datetime import UTC, datetime - -from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from dify_agent.protocol import RunStatusResponse - -from clients.agent_backend import ( - AgentBackendError, - AgentBackendSessionCleanupPayload, - FakeAgentBackendRunClient, - RuntimeLayerSpec, - cleanup_agent_backend_session, -) - - -def _payload() -> AgentBackendSessionCleanupPayload: - return AgentBackendSessionCleanupPayload( - session_snapshot=CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="history", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ) - ] - ), - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - idempotency_key="cleanup-1", - metadata={"tenant_id": "tenant-1"}, - timeout_seconds=15.0, - ) - - -def test_cleanup_agent_backend_session_runs_create_and_wait_until_success(): - client = FakeAgentBackendRunClient(run_id="cleanup-run-1") - - result = cleanup_agent_backend_session(payload=_payload(), client=client) - - assert result.status == "succeeded" - assert result.cleanup_run_id == "cleanup-run-1" - assert client.request is not None - assert [layer.name for layer in client.request.composition.layers] == ["history"] - - -def test_cleanup_agent_backend_session_skips_when_client_is_missing(): - result = cleanup_agent_backend_session( - payload=_payload(), - client=None, - ) - - assert result.status == "skipped" - assert result.reason == "no_agent_backend_client" - - -def test_cleanup_agent_backend_session_skips_when_session_snapshot_is_missing(): - payload = _payload().model_copy(update={"session_snapshot": None}) - - result = cleanup_agent_backend_session(payload=payload, client=FakeAgentBackendRunClient()) - - assert result.status == "skipped" - assert result.reason == "missing_session_snapshot" - - -def test_cleanup_agent_backend_session_skips_when_runtime_layer_specs_are_missing(): - payload = _payload().model_copy(update={"runtime_layer_specs": []}) - - result = cleanup_agent_backend_session(payload=payload, client=FakeAgentBackendRunClient()) - - assert result.status == "skipped" - assert result.reason == "missing_runtime_layer_specs" - - -class _FailedStatusClient(FakeAgentBackendRunClient): - def wait_run(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse: - del timeout_seconds - return RunStatusResponse( - run_id=run_id, - status="failed", - created_at=datetime(2026, 1, 1, tzinfo=UTC), - updated_at=datetime(2026, 1, 1, tzinfo=UTC), - error="snapshot mismatch", - ) - - -def test_cleanup_agent_backend_session_reports_failed_terminal_status(): - client = _FailedStatusClient(run_id="cleanup-run-2") - - result = cleanup_agent_backend_session(payload=_payload(), client=client) - - assert result.status == "failed" - assert result.reason == "snapshot mismatch" - assert result.cleanup_run_id == "cleanup-run-2" - - -class _CreateRunFailureClient(FakeAgentBackendRunClient): - def create_run(self, request): # type: ignore[override] - del request - raise AgentBackendError("create run failed") - - -class _WaitRunFailureClient(FakeAgentBackendRunClient): - def wait_run(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse: - del run_id, timeout_seconds - raise AgentBackendError("wait run failed") - - -def test_cleanup_agent_backend_session_returns_failed_when_create_run_raises(): - result = cleanup_agent_backend_session(payload=_payload(), client=_CreateRunFailureClient()) - - assert result.status == "failed" - assert result.reason == "create run failed" - assert result.cleanup_run_id is None - - -def test_cleanup_agent_backend_session_returns_failed_with_cleanup_run_id_when_wait_run_raises(): - result = cleanup_agent_backend_session( - payload=_payload(), - client=_WaitRunFailureClient(run_id="cleanup-run-3"), - ) - - assert result.status == "failed" - assert result.reason == "wait run failed" - assert result.cleanup_run_id == "cleanup-run-3" diff --git a/api/tests/unit_tests/controllers/common/test_agent_app_parameters.py b/api/tests/unit_tests/controllers/common/test_agent_app_parameters.py index 85d79606b7f..80ca42b0ddc 100644 --- a/api/tests/unit_tests/controllers/common/test_agent_app_parameters.py +++ b/api/tests/unit_tests/controllers/common/test_agent_app_parameters.py @@ -61,6 +61,7 @@ def _persist_snapshot( tenant_id=tenant_id, agent_id=agent_id, version=1, + home_snapshot_id=_stable_uuid(f"home-snapshot:{snapshot_id}"), config_snapshot=config_snapshot, ) session.add(snapshot) diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py index 60c33073408..a56f33ffff2 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py @@ -279,12 +279,7 @@ def test_agent_app_list_and_create_use_agent_route( ) monkeypatch.setattr( roster_controller.AgentRosterService, - "get_or_create_agent_app_debug_conversation_id", - lambda _self, **kwargs: "debug-conversation-detail", - ) - monkeypatch.setattr( - roster_controller.AgentRosterService, - "get_or_create_agent_app_debug_conversation_id", + "get_or_create_build_conversation", lambda _self, **kwargs: "debug-conversation-detail", ) monkeypatch.setattr( @@ -314,7 +309,7 @@ def test_agent_app_list_and_create_use_agent_route( ) monkeypatch.setattr( roster_controller.AgentRosterService, - "load_or_create_agent_app_debug_conversation_ids_by_agent_id", + "load_or_create_build_conversation_ids_by_agent_id", lambda _self, **kwargs: {"agent-list": "debug-conversation-list"}, ) monkeypatch.setattr( @@ -327,7 +322,7 @@ def test_agent_app_list_and_create_use_agent_route( monkeypatch.setattr( roster_controller.AgentRosterService, - "get_or_create_agent_app_debug_conversation_id", + "get_or_create_build_conversation", get_or_create_debug_conversation, ) monkeypatch.setattr( @@ -387,7 +382,6 @@ def test_agent_app_list_and_create_use_agent_route( "tenant_id": "tenant-1", "agent_id": "agent-created", "account_id": account_id, - "draft_type": AgentConfigDraftType.DEBUG_BUILD, "commit": False, } @@ -454,7 +448,7 @@ def test_agent_app_detail_update_delete_resolve_app_from_agent_id( monkeypatch.setattr(roster_controller.AgentRosterService, "get_app_backing_agent", lambda _self, **kwargs: agent) monkeypatch.setattr( roster_controller.AgentRosterService, - "get_or_create_agent_app_debug_conversation_id", + "get_or_create_build_conversation", lambda _self, **kwargs: "debug-conversation-detail", ) monkeypatch.setattr( @@ -561,25 +555,16 @@ def test_agent_app_copy_uses_agent_id_and_returns_agent_detail( } -@pytest.mark.parametrize( - ("payload", "expected_draft_type"), - [ - (None, AgentConfigDraftType.DEBUG_BUILD), - ({"draft_type": "draft"}, AgentConfigDraftType.DRAFT), - ], -) -def test_agent_debug_conversation_refresh_uses_current_user_and_draft_type( +def test_agent_debug_conversation_refresh_resets_build_for_current_user( app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str, - payload: dict[str, str] | None, - expected_draft_type: AgentConfigDraftType, ) -> None: agent_id = "00000000-0000-0000-0000-000000000001" captured: dict[str, object] = {} class FakeRosterService: - def refresh_agent_app_debug_conversation_id(self, **kwargs: object) -> str: + def reset_build_conversation(self, **kwargs: object) -> str: captured.update(kwargs) return "new-debug-conversation-id" @@ -587,7 +572,6 @@ def test_agent_debug_conversation_refresh_uses_current_user_and_draft_type( with app.test_request_context( "/console/api/agent/00000000-0000-0000-0000-000000000001/debug-conversation/refresh", method="POST", - json=payload, ): response = unwrap(AgentDebugConversationRefreshApi.post)( AgentDebugConversationRefreshApi(), MagicMock(), "tenant-1", SimpleNamespace(id=account_id), agent_id @@ -601,7 +585,6 @@ def test_agent_debug_conversation_refresh_uses_current_user_and_draft_type( "tenant_id": "tenant-1", "agent_id": agent_id, "account_id": account_id, - "draft_type": expected_draft_type, } @@ -840,7 +823,7 @@ def test_agent_app_update_allows_empty_role(app: Flask, monkeypatch: pytest.Monk ) monkeypatch.setattr( roster_controller.AgentRosterService, - "get_or_create_agent_app_debug_conversation_id", + "get_or_create_build_conversation", lambda _self, **kwargs: "debug-conversation-detail", ) monkeypatch.setattr( @@ -1203,6 +1186,10 @@ def test_workflow_composer_get_uses_write_transaction() -> None: assert "@with_session\n @get_app_model" in getsource(WorkflowAgentComposerApi) +def test_build_draft_apply_leaves_transaction_ownership_to_service() -> None: + assert "@with_session(write=False)\n def post" in getsource(AgentBuildDraftApplyApi) + + def test_workflow_composer_copy_from_roster(app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str) -> None: app_model = SimpleNamespace(id="app-1") captured: dict[str, object] = {} @@ -1500,7 +1487,7 @@ def test_build_chat_finalization_helper_forces_debug_build_and_push_prompt( assert args["conversation_id"] == "debug-conversation-1" assert args["inputs"] == {} assert args["auto_generate_name"] is False - assert args[completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG] == "delete" + assert "_agent_runtime_exit_intent" not in args assert args["external_trace_id"] == "trace-1" @@ -1643,7 +1630,7 @@ def test_agent_chat_helper_ignores_private_exit_intent_payload_key( "inputs": {}, "query": "hello", "response_mode": "streaming", - completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG: "delete", + "_agent_runtime_exit_intent": "delete", } ): result = completion_controller._create_chat_message( @@ -1657,7 +1644,7 @@ def test_agent_chat_helper_ignores_private_exit_intent_payload_key( args = cast(dict[str, object], captured["args"]) assert args["response_mode"] == "streaming" assert args["conversation_id"] == "debug-conversation-1" - assert completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG not in args + assert "_agent_runtime_exit_intent" not in args @pytest.mark.parametrize( @@ -1717,14 +1704,18 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app def __init__(self, session: object) -> None: calls.append({"session": session}) - def get_or_create_agent_app_debug_conversation_id(self, **kwargs: object) -> str: - calls.append({"get_or_create": kwargs}) + def get_or_create_build_conversation(self, **kwargs: object) -> str: + calls.append({"get_build": kwargs}) return f"debug-{kwargs['agent_id']}" - def refresh_agent_app_debug_conversation_id(self, **kwargs: object) -> str: - calls.append({"refresh": kwargs}) + def rotate_preview_conversation(self, **kwargs: object) -> str: + calls.append({"rotate_preview": kwargs}) return f"new-{kwargs['agent_id']}" + def get_current_preview_conversation(self, **kwargs: object) -> str: + calls.append({"get_preview": kwargs}) + return f"preview-{kwargs['agent_id']}" + def get_app_backing_agent(self, **kwargs: object) -> object: calls.append({"get_app_backing_agent": kwargs}) return SimpleNamespace(id="backing-agent") @@ -1756,33 +1747,46 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app draft_type=AgentConfigDraftType.DRAFT, start_new=True, ) + current_preview_id = completion_controller._resolve_current_user_agent_debug_conversation_id( + session="session-1", # type: ignore[arg-type] + current_tenant_id="tenant-1", + current_user=SimpleNamespace(id="account-1"), + app_model=SimpleNamespace(id="app-1"), + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + ) assert explicit_id == "new-agent-1" assert fallback_id == "debug-backing-agent" assert fallback_preview_id == "new-backing-agent" + assert current_preview_id == "preview-agent-1" assert calls[1] == { - "refresh": { + "rotate_preview": { "tenant_id": "tenant-1", "agent_id": "agent-1", "account_id": "account-1", - "draft_type": AgentConfigDraftType.DRAFT, } } assert calls[3] == {"get_app_backing_agent": {"tenant_id": "tenant-1", "app_id": "app-1"}} assert calls[4] == { - "get_or_create": { + "get_build": { "tenant_id": "tenant-1", "agent_id": "backing-agent", "account_id": "account-1", - "draft_type": AgentConfigDraftType.DEBUG_BUILD, } } assert calls[6] == {"get_app_backing_agent": {"tenant_id": "tenant-1", "app_id": "app-1"}} assert calls[7] == { - "refresh": { + "rotate_preview": { "tenant_id": "tenant-1", "agent_id": "backing-agent", "account_id": "account-1", - "draft_type": AgentConfigDraftType.DRAFT, + } + } + assert calls[9] == { + "get_preview": { + "tenant_id": "tenant-1", + "agent_id": "agent-1", + "account_id": "account-1", } } diff --git a/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py b/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py index 265c542e635..1a1b6f87916 100644 --- a/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py +++ b/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock import pytest from dify_agent.client import DifyAgentClientError, DifyAgentHTTPError, DifyAgentTimeoutError -from dify_agent.protocol import SandboxListResponse, SandboxReadResponse +from dify_agent.protocol import WorkspaceListResponse, WorkspaceReadResponse from controllers.console import agent_app_sandbox as module from models.model import App, AppMode, IconType @@ -15,30 +15,67 @@ from services.agent_app_sandbox_service import AgentSandboxInfo, AgentSandboxIns class _AgentAppService: def __init__(self) -> None: - self.calls: list[tuple[str, str, str, str, str]] = [] + self.calls: list[tuple[str, str, str, str, str, str, str, str]] = [] - def get_info(self, *, tenant_id: str, app_id: str, conversation_id: str) -> AgentSandboxInfo: - self.calls.append(("info", tenant_id, app_id, conversation_id, "")) - return AgentSandboxInfo(session_id="abc1234", workspace_cwd="~/workspace/abc1234") + def get_info( + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: str, + caller_id: str, + account_id: str, + ) -> AgentSandboxInfo: + self.calls.append(("info", tenant_id, app_id, agent_id, caller_type, caller_id, account_id, "")) + return AgentSandboxInfo(workspace_cwd=".") - def list_files(self, *, tenant_id: str, app_id: str, conversation_id: str, path: str) -> SandboxListResponse: - self.calls.append(("list", tenant_id, app_id, conversation_id, path)) - return SandboxListResponse(path=path, entries=[], truncated=False) + def list_files( + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: str, + caller_id: str, + account_id: str, + path: str, + ) -> WorkspaceListResponse: + self.calls.append(("list", tenant_id, app_id, agent_id, caller_type, caller_id, account_id, path)) + return WorkspaceListResponse(path=path, entries=[], truncated=False) - def read_file(self, *, tenant_id: str, app_id: str, conversation_id: str, path: str) -> SandboxReadResponse: - self.calls.append(("read", tenant_id, app_id, conversation_id, path)) - return SandboxReadResponse(path=path, size=5, truncated=False, binary=False, text="hello") + def read_file( + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: str, + caller_id: str, + account_id: str, + path: str, + ) -> WorkspaceReadResponse: + self.calls.append(("read", tenant_id, app_id, agent_id, caller_type, caller_id, account_id, path)) + return WorkspaceReadResponse(path=path, size=5, truncated=False, binary=False, text="hello") def upload_file( - self, *, tenant_id: str, app_id: str, conversation_id: str, path: str + self, + *, + tenant_id: str, + app_id: str, + agent_id: str, + caller_type: str, + caller_id: str, + account_id: str, + path: str, ) -> AgentSandboxUploadDownload: - self.calls.append(("upload", tenant_id, app_id, conversation_id, path)) + self.calls.append(("upload", tenant_id, app_id, agent_id, caller_type, caller_id, account_id, path)) return AgentSandboxUploadDownload(url="https://files.example/report.txt") class _WorkflowService: def __init__(self) -> None: - self.calls: list[tuple[str, str, str, str, str, str | None, str]] = [] + self.calls: list[tuple[str, str, str, str, str, str]] = [] def list_files( self, @@ -47,12 +84,12 @@ class _WorkflowService: app_id: str, workflow_run_id: str, node_id: str, - node_execution_id: str | None, + node_execution_id: str, path: str, session, - ) -> SandboxListResponse: - self.calls.append(("list", tenant_id, app_id, workflow_run_id, node_id, node_execution_id, path)) - return SandboxListResponse(path=path, entries=[], truncated=False) + ) -> WorkspaceListResponse: + self.calls.append(("list", tenant_id, app_id, workflow_run_id, node_id, path)) + return WorkspaceListResponse(path=path, entries=[], truncated=False) def read_file( self, @@ -61,12 +98,12 @@ class _WorkflowService: app_id: str, workflow_run_id: str, node_id: str, - node_execution_id: str | None, + node_execution_id: str, path: str, session, - ) -> SandboxReadResponse: - self.calls.append(("read", tenant_id, app_id, workflow_run_id, node_id, node_execution_id, path)) - return SandboxReadResponse(path=path, size=5, truncated=False, binary=False, text="hello") + ) -> WorkspaceReadResponse: + self.calls.append(("read", tenant_id, app_id, workflow_run_id, node_id, path)) + return WorkspaceReadResponse(path=path, size=5, truncated=False, binary=False, text="hello") def upload_file( self, @@ -75,11 +112,11 @@ class _WorkflowService: app_id: str, workflow_run_id: str, node_id: str, - node_execution_id: str | None, + node_execution_id: str, path: str, session, ) -> AgentSandboxUploadDownload: - self.calls.append(("upload", tenant_id, app_id, workflow_run_id, node_id, node_execution_id, path)) + self.calls.append(("upload", tenant_id, app_id, workflow_run_id, node_id, path)) return AgentSandboxUploadDownload(url="https://files.example/upload.txt") @@ -125,34 +162,41 @@ def test_handle_maps_sandbox_and_agent_backend_errors() -> None: def test_agent_app_sandbox_resources_proxy_service(monkeypatch: pytest.MonkeyPatch) -> None: service = _AgentAppService() session = MagicMock() + account = SimpleNamespace(id="account-1") resolver = MagicMock(return_value=_app_model()) monkeypatch.setattr(module, "AgentAppSandboxService", lambda: service) monkeypatch.setattr(module, "resolve_agent_runtime_app_model", resolver) monkeypatch.setattr( module, "query_params_from_request", - lambda model: SimpleNamespace(conversation_id="conv-1", path="sub/report.txt"), + lambda model: SimpleNamespace(caller_type="build_draft", caller_id="build-1", path="sub/report.txt"), ) monkeypatch.setattr( module, "request", - SimpleNamespace(get_json=lambda silent=True: {"conversation_id": "conv-1", "path": "report.txt"}), + SimpleNamespace( + get_json=lambda silent=True: { + "caller_type": "build_draft", + "caller_id": "build-1", + "path": "report.txt", + } + ), ) - info = unwrap(module.AgentAppSandboxInfoResource.get)(object(), session, "tenant-1", "agent-1") - listing = unwrap(module.AgentAppSandboxListResource.get)(object(), session, "tenant-1", "agent-1") - preview = unwrap(module.AgentAppSandboxReadResource.get)(object(), session, "tenant-1", "agent-1") - upload = unwrap(module.AgentAppSandboxUploadResource.post)(object(), session, "tenant-1", "agent-1") + info = unwrap(module.AgentAppSandboxInfoResource.get)(object(), session, account, "tenant-1", "agent-1") + listing = unwrap(module.AgentAppSandboxListResource.get)(object(), session, account, "tenant-1", "agent-1") + preview = unwrap(module.AgentAppSandboxReadResource.get)(object(), session, account, "tenant-1", "agent-1") + upload = unwrap(module.AgentAppSandboxUploadResource.post)(object(), session, account, "tenant-1", "agent-1") - assert info == {"session_id": "abc1234", "workspace_cwd": "~/workspace/abc1234"} + assert info == {"workspace_cwd": "."} assert listing["path"] == "sub/report.txt" assert preview["text"] == "hello" assert upload == {"url": "https://files.example/report.txt"} assert service.calls == [ - ("info", "tenant-1", "app-1", "conv-1", ""), - ("list", "tenant-1", "app-1", "conv-1", "sub/report.txt"), - ("read", "tenant-1", "app-1", "conv-1", "sub/report.txt"), - ("upload", "tenant-1", "app-1", "conv-1", "report.txt"), + ("info", "tenant-1", "app-1", "agent-1", "build_draft", "build-1", "account-1", ""), + ("list", "tenant-1", "app-1", "agent-1", "build_draft", "build-1", "account-1", "sub/report.txt"), + ("read", "tenant-1", "app-1", "agent-1", "build_draft", "build-1", "account-1", "sub/report.txt"), + ("upload", "tenant-1", "app-1", "agent-1", "build_draft", "build-1", "account-1", "report.txt"), ] assert all(call.kwargs["session"] is session for call in resolver.call_args_list) @@ -160,24 +204,27 @@ def test_agent_app_sandbox_resources_proxy_service(monkeypatch: pytest.MonkeyPat def test_agent_app_sandbox_resource_returns_normalized_errors(monkeypatch: pytest.MonkeyPatch) -> None: class FailingService: def get_info(self, **kwargs): - raise AgentSandboxInspectorError("no_active_session", "no active session", status_code=404) + raise AgentSandboxInspectorError("no_active_binding", "no active binding", status_code=404) def list_files(self, **kwargs): - raise AgentSandboxInspectorError("no_active_session", "no active session", status_code=404) + raise AgentSandboxInspectorError("no_active_binding", "no active binding", status_code=404) monkeypatch.setattr(module, "AgentAppSandboxService", FailingService) session = MagicMock() + account = SimpleNamespace(id="account-1") monkeypatch.setattr(module, "resolve_agent_runtime_app_model", MagicMock(return_value=_app_model())) monkeypatch.setattr( - module, "query_params_from_request", lambda model: SimpleNamespace(conversation_id="conv-1", path=".") + module, + "query_params_from_request", + lambda model: SimpleNamespace(caller_type="conversation", caller_id="conv-1", path="."), ) - assert unwrap(module.AgentAppSandboxInfoResource.get)(object(), session, "tenant-1", "agent-1") == ( - {"code": "no_active_session", "message": "no active session"}, + assert unwrap(module.AgentAppSandboxInfoResource.get)(object(), session, account, "tenant-1", "agent-1") == ( + {"code": "no_active_binding", "message": "no active binding"}, 404, ) - assert unwrap(module.AgentAppSandboxListResource.get)(object(), session, "tenant-1", "agent-1") == ( - {"code": "no_active_session", "message": "no active session"}, + assert unwrap(module.AgentAppSandboxListResource.get)(object(), session, account, "tenant-1", "agent-1") == ( + {"code": "no_active_binding", "message": "no active binding"}, 404, ) @@ -188,12 +235,12 @@ def test_workflow_agent_sandbox_resources_proxy_service(monkeypatch: pytest.Monk monkeypatch.setattr( module, "query_params_from_request", - lambda model: SimpleNamespace(node_execution_id="exec-1", path="out.txt"), + lambda model: SimpleNamespace(node_execution_id="execution-1", path="out.txt"), ) monkeypatch.setattr( module, "request", - SimpleNamespace(get_json=lambda silent=True: {"node_execution_id": "exec-1", "path": "upload.txt"}), + SimpleNamespace(get_json=lambda silent=True: {"node_execution_id": "execution-1", "path": "upload.txt"}), ) app_model = _app_model() @@ -211,7 +258,7 @@ def test_workflow_agent_sandbox_resources_proxy_service(monkeypatch: pytest.Monk assert preview["text"] == "hello" assert upload == {"url": "https://files.example/upload.txt"} assert service.calls == [ - ("list", "tenant-1", "app-1", "run-1", "agent-node", "exec-1", "out.txt"), - ("read", "tenant-1", "app-1", "run-1", "agent-node", "exec-1", "out.txt"), - ("upload", "tenant-1", "app-1", "run-1", "agent-node", "exec-1", "upload.txt"), + ("list", "tenant-1", "app-1", "run-1", "agent-node", "out.txt"), + ("read", "tenant-1", "app-1", "run-1", "agent-node", "out.txt"), + ("upload", "tenant-1", "app-1", "run-1", "agent-node", "upload.txt"), ] diff --git a/api/tests/unit_tests/controllers/test_swagger.py b/api/tests/unit_tests/controllers/test_swagger.py index 149af9ff76a..3b0bb6e4e63 100644 --- a/api/tests/unit_tests/controllers/test_swagger.py +++ b/api/tests/unit_tests/controllers/test_swagger.py @@ -574,7 +574,7 @@ def test_console_account_avatar_query_param_renders_as_query(monkeypatch: pytest assert params["avatar"]["required"] is True -def test_console_agent_debug_conversation_refresh_body_is_optional(monkeypatch: pytest.MonkeyPatch): +def test_console_agent_debug_conversation_refresh_has_no_body(monkeypatch: pytest.MonkeyPatch): from configs import dify_config from controllers.console import bp as console_bp @@ -586,13 +586,9 @@ def test_console_agent_debug_conversation_refresh_body_is_optional(monkeypatch: payload = app.test_client().get("/console/api/openapi.json").get_json() operation = payload["paths"]["/agent/{agent_id}/debug-conversation/refresh"]["post"] - request_body = operation["requestBody"] - assert request_body["required"] is False - assert request_body["content"]["application/json"]["schema"] == { - "$ref": "#/components/schemas/AgentDebugConversationRefreshPayload" - } - assert "AgentDebugConversationRefreshPayload" in payload["components"]["schemas"] + assert "requestBody" not in operation + assert "AgentDebugConversationRefreshPayload" not in payload["components"]["schemas"] def test_console_member_invite_documents_bad_request_response(monkeypatch: pytest.MonkeyPatch): diff --git a/api/tests/unit_tests/core/agent/test_publish_visibility.py b/api/tests/unit_tests/core/agent/test_publish_visibility.py index 7e6b0964c83..042c88f3d17 100644 --- a/api/tests/unit_tests/core/agent/test_publish_visibility.py +++ b/api/tests/unit_tests/core/agent/test_publish_visibility.py @@ -71,6 +71,7 @@ def _add_agent( tenant_id="tenant-1", agent_id=agent_id, version=1, + home_snapshot_id=f"home-{agent_id}", config_snapshot=_agent_soul(), ) ) @@ -245,6 +246,7 @@ def test_workflow_callable_filter_distinguishes_never_published_from_dirty_draft tenant_id="tenant-1", agent_id=stale_published_snapshot.id, version=2, + home_snapshot_id="home-stale-published-old", config_snapshot=_agent_soul(), ) ) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py index 12c429e85b9..d193231b96e 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py @@ -22,11 +22,10 @@ from core.app.apps.agent_app.app_generator import ( AgentAppGeneratorError, ) from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import AGENT_RUNTIME_EXIT_INTENT_ARG, InvokeFrom, UserFrom +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.app.entities.queue_entities import QueueAnnotationReplyEvent from core.workflow.file_reference import build_file_reference from models import Account, AppModelConfig -from models.agent import AgentConfigDraftType MODULE = "core.app.apps.agent_app.app_generator" @@ -79,14 +78,18 @@ class TestGenerateGuards: class TestGenerateSuccess: - def test_runtime_session_snapshot_id_preserves_snapshot_for_debugger_and_web_app(self): + def test_session_scope_config_version_id_preserves_draft_or_snapshot_id(self): assert ( - AgentAppGenerator._runtime_session_snapshot_id(invoke_from=InvokeFrom.DEBUGGER, snapshot_id="snap-1") - == "snap-1" + AgentAppGenerator._session_scope_config_version_id( + invoke_from=InvokeFrom.DEBUGGER, config_version_id="draft-1" + ) + == "draft-1" ) assert ( - AgentAppGenerator._runtime_session_snapshot_id(invoke_from=InvokeFrom.WEB_APP, snapshot_id="snap-1") - == "snap-1" + AgentAppGenerator._session_scope_config_version_id( + invoke_from=InvokeFrom.WEB_APP, config_version_id="snapshot-1" + ) + == "snapshot-1" ) def test_generate_orchestrates_and_starts_worker(self, generator, mocker: MockerFixture): @@ -145,84 +148,11 @@ class TestGenerateSuccess: draft_type=None, user=user, session=session, + conversation=None, ) session.get.assert_called_once_with(AppModelConfig, "config-1") assert generate_entity.call_args.kwargs["prompt_file_mappings"] == file_mappings - assert generate_entity.call_args.kwargs["agent_runtime_exit_intent"] == "suspend" - - def test_generate_uses_delete_exit_intent_from_internal_arg(self, generator, mocker: MockerFixture): - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") - user = DummyAccount("user") - - generator._resolve_agent = mocker.MagicMock( - return_value=(mocker.MagicMock(id="agent1"), "snap1", "snapshot", mocker.MagicMock()) - ) - generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock( - return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) - ) - generator._handle_response = mocker.MagicMock(return_value="raw-response") - - mocker.patch( - f"{MODULE}.AgentAppConfigManager.get_app_config", - return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), - ) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - generate_entity = mocker.patch( - f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") - ) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - - generator.generate( - app_model=app_model, - user=user, - args={"query": "hello", "inputs": {}, AGENT_RUNTIME_EXIT_INTENT_ARG: "delete"}, - invoke_from=InvokeFrom.DEBUGGER, - session=mocker.MagicMock(), - streaming=True, - ) - - assert generate_entity.call_args.kwargs["agent_runtime_exit_intent"] == "delete" - - def test_generate_falls_back_to_suspend_for_invalid_internal_exit_intent(self, generator, mocker: MockerFixture): - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") - user = DummyAccount("user") - - generator._resolve_agent = mocker.MagicMock( - return_value=(mocker.MagicMock(id="agent1"), "snap1", "snapshot", mocker.MagicMock()) - ) - generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock( - return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) - ) - generator._handle_response = mocker.MagicMock(return_value="raw-response") - - mocker.patch( - f"{MODULE}.AgentAppConfigManager.get_app_config", - return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), - ) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - generate_entity = mocker.patch( - f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") - ) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - - generator.generate( - app_model=app_model, - user=user, - args={"query": "hello", "inputs": {}, AGENT_RUNTIME_EXIT_INTENT_ARG: "bogus"}, - invoke_from=InvokeFrom.DEBUGGER, - session=mocker.MagicMock(), - streaming=True, - ) - - assert generate_entity.call_args.kwargs["agent_runtime_exit_intent"] == "suspend" + assert "agent_runtime_exit_intent" not in generate_entity.call_args.kwargs def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") @@ -263,6 +193,7 @@ class TestGenerateSuccess: user=user, session=session, ) + assert generator._resolve_agent.call_args.kwargs["conversation"].id == "conv" assert generator._init_generate_records.call_args.kwargs["session"] is session def test_generate_does_not_include_trace_session_id_in_extras( @@ -326,8 +257,10 @@ class TestGenerateWorker: generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock(id="conv")) generator._get_message = mocker.MagicMock(return_value=mocker.MagicMock(id="msg")) generator._run_input_guards = mocker.MagicMock(return_value=(handled, guard_query, None)) + resolved_agent = mocker.MagicMock(id="a") + resolved_config = mocker.MagicMock(id="s", home_snapshot_id="home-1") generator._resolve_agent_by_id = mocker.MagicMock( - return_value=(mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()) + return_value=(resolved_agent, resolved_config, mocker.MagicMock()) ) session = mocker.MagicMock() session.get.return_value = mocker.MagicMock(id="app1") @@ -345,7 +278,7 @@ class TestGenerateWorker: mocker.patch(f"{MODULE}.AgentAppRuntimeRequestBuilder", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.create_agent_backend_run_client", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.AgentBackendRunEventAdapter", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppRuntimeSessionStore", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppWorkspaceStore", return_value=mocker.MagicMock()) runner = mocker.MagicMock() if run_side_effect is not None: runner.run.side_effect = run_side_effect @@ -360,9 +293,8 @@ class TestGenerateWorker: *, is_resume=False, query="query", - runtime_session_snapshot_id="s", + session_scope_config_version_id="s", prompt_file_mappings=(), - agent_runtime_exit_intent="suspend", ): generator._generate_worker( flask_app=mocker.MagicMock(), @@ -370,8 +302,7 @@ class TestGenerateWorker: application_generate_entity=mocker.MagicMock( agent_id="a", agent_config_snapshot_id="s", - agent_runtime_session_snapshot_id=runtime_session_snapshot_id, - agent_runtime_exit_intent=agent_runtime_exit_intent, + agent_session_scope_config_version_id=session_scope_config_version_id, model_conf=mocker.MagicMock(model="m"), query=query, prompt_file_mappings=prompt_file_mappings, @@ -389,25 +320,19 @@ class TestGenerateWorker: self._call(generator, mocker, queue_manager) runner.run.assert_called_once() assert generator._resolve_agent_by_id.call_args.kwargs["session"] is resolver_session + assert runner.run.call_args.kwargs["home_snapshot_id"] == "home-1" + assert "home_snapshot_ref" not in runner.run.call_args.kwargs queue_manager.publish_error.assert_not_called() - def test_worker_passes_runtime_session_scope_to_runner(self, generator, mocker: MockerFixture): + def test_worker_passes_session_scope_config_version_to_runner(self, generator, mocker: MockerFixture): runner, _ = self._wire(generator, mocker) queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager, runtime_session_snapshot_id=None) + self._call(generator, mocker, queue_manager, session_scope_config_version_id=None) assert runner.run.call_args.kwargs["agent_config_snapshot_id"] == "s" assert runner.run.call_args.kwargs["session_scope_snapshot_id"] is None - def test_worker_forwards_runtime_exit_intent_to_runner(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker) - queue_manager = mocker.MagicMock() - - self._call(generator, mocker, queue_manager, agent_runtime_exit_intent="delete") - - assert runner.run.call_args.kwargs["agent_runtime_exit_intent"] == "delete" - def test_worker_appends_prompt_files_to_backend_query(self, generator, mocker: MockerFixture): runner, _ = self._wire(generator, mocker, guard_query="你看得见这张图片吗") queue_manager = mocker.MagicMock() @@ -526,7 +451,7 @@ class TestResumeAfterFormSubmission: mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppRuntimeSessionStore") + generator._resolve_resume_draft = mocker.MagicMock(return_value=(None, None)) return ( mocker.patch( f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") @@ -547,6 +472,7 @@ class TestResumeAfterFormSubmission: app_model=app_model, user=user, conversation_id="conv", + form_id="form-1", invoke_from=InvokeFrom.WEB_APP, session=session, ) @@ -573,6 +499,7 @@ class TestResumeAfterFormSubmission: app_model=mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent"), user=DummyAccount("user"), conversation_id="conv", + form_id="form-1", invoke_from=InvokeFrom.WEB_APP, session=session, ) @@ -584,25 +511,24 @@ class TestResumeAfterFormSubmission: self._wire(generator, mocker) conversation = mocker.MagicMock(id="conv", invoke_from=InvokeFrom.DEBUGGER) mocker.patch(f"{MODULE}.ConversationService.get_conversation", return_value=conversation) - session_store = mocker.patch(f"{MODULE}.AgentAppRuntimeSessionStore") - session_store.return_value.load_active_session_for_conversation.return_value = mocker.MagicMock( - scope=mocker.MagicMock(agent_config_snapshot_id="draft-build-1") - ) - draft_row = mocker.MagicMock(draft_type=AgentConfigDraftType.DEBUG_BUILD, account_id="user") + generator._resolve_resume_draft.return_value = ("debug_build", "draft-build-1") account_user = mocker.MagicMock(spec=Account) account_user.id = "user" app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") app_model.app_model_config_id = "config-1" session = mocker.MagicMock() - session.scalar.side_effect = [draft_row, mocker.MagicMock(query="original question")] + session.scalar.return_value = mocker.MagicMock(query="original question") generator.resume_after_form_submission( app_model=app_model, user=account_user, conversation_id="conv", + form_id="form-1", invoke_from=InvokeFrom.DEBUGGER, session=session, ) assert generator._resolve_agent.call_args.kwargs["draft_type"] == "debug_build" + assert generator._resolve_agent.call_args.kwargs["draft_id"] == "draft-build-1" assert generator._resolve_agent.call_args.kwargs["session"] is session + assert generator._resolve_agent.call_args.kwargs["conversation"] is conversation diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py index f623dd767d2..95505a9949f 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py @@ -23,7 +23,6 @@ from dify_agent.protocol import ( RunStartedEvent, RunSucceededEvent, RunSucceededEventData, - RuntimeLayerSpec, ) from pydantic_ai.messages import ( FunctionToolCallEvent, @@ -36,7 +35,6 @@ from pydantic_ai.messages import ( ) from clients.agent_backend import ( - AgentBackendError, AgentBackendRunEventAdapter, AgentBackendRunFailedError, AgentBackendRunFailedInternalEvent, @@ -155,43 +153,6 @@ class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient): ) -class _StreamingRecordingFakeAgentBackendRunClient(_RecordingFakeAgentBackendRunClient): - @override - def stream_events( - self, - run_id: str, - *, - after: str | None = None, - should_stop: Callable[[], bool] | None = None, - ) -> Iterator[RunEvent]: - del after, should_stop - created_at = datetime(2026, 1, 1, tzinfo=UTC) - yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at) - yield PydanticAIStreamRunEvent( - id="2-0", - run_id=run_id, - created_at=created_at, - data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hello ")), - agent_message_delta="hello ", - ) - yield PydanticAIStreamRunEvent( - id="3-0", - run_id=run_id, - created_at=created_at, - data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="agent")), - agent_message_delta="agent", - ) - yield RunSucceededEvent( - id="4-0", - run_id=run_id, - created_at=created_at, - data=RunSucceededEventData( - output={"text": "hello agent"}, - session_snapshot=CompositorSessionSnapshot(layers=[]), - ), - ) - - class _StreamingStopAfterFirstDeltaFakeAgentBackendRunClient(_RecordingFakeAgentBackendRunClient): def __init__(self, *, queue_manager: _FakeQueueManager, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -419,60 +380,48 @@ class _FakeSessionStore: self, loaded: CompositorSessionSnapshot | None = None, loaded_session: StoredAgentAppSession | None = None, - listed_sessions: list[StoredAgentAppSession] | None = None, + binding_id: str = "binding-1", + workspace_id: str = "workspace-1", + backend_binding_ref: str = "backend-binding-1", ) -> None: self.loaded = loaded self._loaded_session = loaded_session - self._listed_sessions = list(listed_sessions or []) - self.loaded_scopes: list[AgentAppSessionScope] = [] + self.binding_id = binding_id + self.workspace_id = workspace_id + self.backend_binding_ref = backend_binding_ref + self.resolved_scopes: list[AgentAppSessionScope] = [] self.saved: list[ tuple[ AgentAppSessionScope, str, CompositorSessionSnapshot | None, - list[RuntimeLayerSpec], str | None, str | None, ] ] = [] - self.cleaned: list[tuple[AgentAppSessionScope, str | None]] = [] - def load_active_snapshot(self, scope: AgentAppSessionScope) -> CompositorSessionSnapshot | None: - self.loaded_scopes.append(scope) - return self.loaded - - def load_active_session(self, scope: AgentAppSessionScope) -> StoredAgentAppSession | None: - self.loaded_scopes.append(scope) + def load_or_create(self, scope: AgentAppSessionScope) -> StoredAgentAppSession: + self.resolved_scopes.append(scope) if self._loaded_session is not None: return self._loaded_session - if self.loaded is None: - return None - return StoredAgentAppSession(scope=scope, session_snapshot=self.loaded, backend_run_id=None) - - def list_active_sessions_for_conversation( - self, *, tenant_id: str, app_id: str, conversation_id: str - ) -> list[StoredAgentAppSession]: - assert tenant_id == "tenant-1" - assert app_id == "app-1" - assert conversation_id == "conv-1" - return list(self._listed_sessions) + return StoredAgentAppSession( + scope=scope, + binding_id=self.binding_id, + workspace_id=self.workspace_id, + backend_binding_ref=self.backend_binding_ref, + session_snapshot=self.loaded, + ) def save_active_snapshot( self, *, - scope, - backend_run_id, - snapshot, - runtime_layer_specs, - pending_form_id=None, - pending_tool_call_id=None, + scope: AgentAppSessionScope, + binding_id: str, + snapshot: CompositorSessionSnapshot | None, + pending_form_id: str | None = None, + pending_tool_call_id: str | None = None, ) -> None: - self.saved.append( - (scope, backend_run_id, snapshot, list(runtime_layer_specs), pending_form_id, pending_tool_call_id) - ) - - def mark_cleaned(self, *, scope: AgentAppSessionScope, backend_run_id: str | None = None) -> None: - self.cleaned.append((scope, backend_run_id)) + self.saved.append((scope, binding_id, snapshot, pending_form_id, pending_tool_call_id)) class _MonotonicClock: @@ -529,18 +478,18 @@ def _runner( ) -def _run(runner: AgentAppRunner, qm: _FakeQueueManager, *, agent_runtime_exit_intent: str = "suspend") -> None: +def _run(runner: AgentAppRunner, qm: _FakeQueueManager) -> None: runner.run( dify_context=_dify_ctx(), agent_id="agent-1", agent_config_snapshot_id="snap-1", agent_soul=_soul(), + home_snapshot_id="home-1", conversation_id="conv-1", query="hello", message_id="msg-1", model_name="gpt-4o-mini", queue_manager=qm, # type: ignore[arg-type] - agent_runtime_exit_intent=agent_runtime_exit_intent, # type: ignore[arg-type] ) @@ -578,134 +527,27 @@ def test_successful_turn_publishes_chunk_and_message_end_and_saves_session(): assert _saved_user_query(qm) == "hello" # The conversation session snapshot is persisted for multi-turn continuity. assert store.saved - saved_scope, saved_run_id, saved_snapshot, saved_specs, pending_form_id, pending_tool_call_id = store.saved[0] + saved_scope, saved_binding_id, saved_snapshot, pending_form_id, pending_tool_call_id = store.saved[0] assert saved_scope.conversation_id == "conv-1" assert saved_scope.agent_config_snapshot_id == "snap-1" - assert saved_run_id == "fake-run-1" + assert saved_binding_id == "binding-1" assert saved_snapshot is not None - assert saved_specs # A successful turn carries no ask_human pause correlation. assert pending_form_id is None assert pending_tool_call_id is None - assert store.cleaned == [] -def test_successful_turn_enqueues_cleanup_for_superseded_sessions_after_saving_snapshot(monkeypatch): - superseded = StoredAgentAppSession( - scope=AgentAppSessionScope( - tenant_id="tenant-1", - app_id="app-1", - conversation_id="conv-1", - agent_id="agent-2", - agent_config_snapshot_id="snap-2", - ), - session_snapshot=CompositorSessionSnapshot(layers=[]), - backend_run_id="run-old", - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - ) - current_scope_session = StoredAgentAppSession( - scope=AgentAppSessionScope( - tenant_id="tenant-1", - app_id="app-1", - conversation_id="conv-1", - agent_id="agent-1", - agent_config_snapshot_id="snap-1", - ), - session_snapshot=CompositorSessionSnapshot(layers=[]), - backend_run_id="run-current", - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - ) - store = _FakeSessionStore(listed_sessions=[current_scope_session, superseded]) +def test_turn_uses_resolved_backend_binding_before_backend_invocation() -> None: client = FakeAgentBackendRunClient() - qm = _FakeQueueManager() - cleanup_delay = MagicMock() - monkeypatch.setattr(app_runner_module.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) + store = _FakeSessionStore(binding_id="binding-2", backend_binding_ref="backend-binding-2") - _run(_runner(client, store), qm) - - assert store.saved - cleanup_delay.assert_called_once() - payload = cleanup_delay.call_args.args[0] - assert payload["metadata"]["conversation_id"] == "conv-1" - assert payload["metadata"]["agent_id"] == "agent-2" - assert payload["metadata"]["previous_agent_backend_run_id"] == "run-old" - assert payload["idempotency_key"] == "tenant-1:app-1:conv-1:agent-2:snap-2:superseded-session-cleanup:run-old" - - -def test_superseded_session_cleanup_enqueue_failure_does_not_fail_turn(monkeypatch): - superseded = StoredAgentAppSession( - scope=AgentAppSessionScope( - tenant_id="tenant-1", - app_id="app-1", - conversation_id="conv-1", - agent_id="agent-2", - agent_config_snapshot_id="snap-2", - ), - session_snapshot=CompositorSessionSnapshot(layers=[]), - backend_run_id="run-old", - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - ) - store = _FakeSessionStore(listed_sessions=[superseded]) - client = FakeAgentBackendRunClient() - qm = _FakeQueueManager() - cleanup_delay = MagicMock(side_effect=RuntimeError("queue down")) - monkeypatch.setattr(app_runner_module.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) - - _run(_runner(client, store), qm) - - cleanup_delay.assert_called_once() - end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert len(end_events) == 1 - assert end_events[0].llm_result.message.content == "hello agent" - - -def test_delete_on_exit_turn_marks_session_cleaned_without_saving_snapshot(): - client = _StreamingRecordingFakeAgentBackendRunClient() - store = _FakeSessionStore() - qm = _FakeQueueManager() - - _run(_runner(client, store), qm, agent_runtime_exit_intent="delete") + _run(_runner(client, store), _FakeQueueManager()) assert client.request is not None - assert client.request.on_exit.default.value == "delete" - assert store.saved == [] - assert len(store.cleaned) == 1 - cleaned_scope, cleaned_run_id = store.cleaned[0] - assert cleaned_scope.conversation_id == "conv-1" - assert cleaned_scope.agent_config_snapshot_id == "snap-1" - assert cleaned_run_id == "fake-run-1" - end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert len(end_events) == 1 - assert end_events[0].llm_result.message.content == "hello agent" - - -def test_delete_on_exit_turn_swallows_cleanup_failure_after_success(): - client = _StreamingRecordingFakeAgentBackendRunClient() - store = _FakeSessionStore() - store.mark_cleaned = MagicMock(side_effect=RuntimeError("cleanup failed")) # type: ignore[method-assign] - qm = _FakeQueueManager() - - _run(_runner(client, store), qm, agent_runtime_exit_intent="delete") - - assert store.saved == [] - store.mark_cleaned.assert_called_once() - end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert len(end_events) == 1 - - -def test_delete_on_exit_turn_marks_session_cleaned_when_publish_fails(): - client = _StreamingRecordingFakeAgentBackendRunClient() - store = _FakeSessionStore() - store.mark_cleaned = MagicMock(side_effect=RuntimeError("cleanup failed")) # type: ignore[method-assign] - qm = _FakeQueueManager() - runner = _runner(client, store) - runner._publish_terminal_answer = MagicMock(side_effect=RuntimeError("publish failed")) - - with pytest.raises(RuntimeError, match="publish failed"): - _run(runner, qm, agent_runtime_exit_intent="delete") - - assert store.saved == [] - store.mark_cleaned.assert_called_once() + layers = {layer["name"]: layer for layer in client.request.model_dump(mode="json")["composition"]["layers"]} + assert layers["runtime"]["config"]["backend_binding_ref"] == "backend-binding-2" + assert store.saved[0][1] == "binding-2" + assert len(store.resolved_scopes) == 1 def test_successful_turn_routes_stream_text_to_agent_message_and_uses_terminal_output(monkeypatch): @@ -1254,6 +1096,7 @@ def test_debug_session_scope_can_reuse_conversation_across_config_snapshots(): agent_id="agent-1", agent_config_snapshot_id="snap-new", agent_soul=_soul(), + home_snapshot_id="home-1", conversation_id="conv-1", query="hello", message_id="msg-1", @@ -1264,8 +1107,8 @@ def test_debug_session_scope_can_reuse_conversation_across_config_snapshots(): assert client.request is not None assert client.request.session_snapshot is prior - assert store.loaded_scopes[0].agent_config_snapshot_id is None - assert store.saved[0][0].agent_config_snapshot_id is None + assert store.resolved_scopes[0].agent_config_snapshot_id == "snap-new" + assert store.saved[0][0].agent_config_snapshot_id == "snap-new" def test_failed_run_raises_agent_backend_error(): @@ -1358,24 +1201,8 @@ def test_ask_human_pauses_turn_creates_form_and_persists_correlation(): assert _saved_user_query(qm) == "hello" # The pause correlation is persisted so a form submission can resume the run. assert store.saved - assert store.saved[0][4] == "form-1" - assert store.saved[0][5] == "fake-ask-human-1" - - -def test_delete_on_exit_deferred_tool_marks_session_cleaned_and_raises_error(): - client = FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario.PAUSED) - store = _FakeSessionStore() - store.mark_cleaned = MagicMock(side_effect=RuntimeError("cleanup failed")) # type: ignore[method-assign] - qm = _FakeQueueManager() - runner = _runner(client, store) - runner._pause_for_ask_human = MagicMock() - - with pytest.raises(AgentBackendError, match="finalization cannot pause for human input"): - _run(runner, qm, agent_runtime_exit_intent="delete") - - runner._pause_for_ask_human.assert_not_called() - assert store.saved == [] - store.mark_cleaned.assert_called_once() + assert store.saved[0][3] == "form-1" + assert store.saved[0][4] == "fake-ask-human-1" def test_submitted_form_resumes_turn_with_deferred_tool_results(monkeypatch): @@ -1389,9 +1216,12 @@ def test_submitted_form_resumes_turn_with_deferred_tool_results(monkeypatch): conversation_id="conv-1", agent_id="agent-1", agent_config_snapshot_id="snap-1", + home_snapshot_id="home-1", ), + binding_id="binding-1", + workspace_id="workspace-1", + backend_binding_ref="backend-binding-1", session_snapshot=snapshot, - backend_run_id="run-0", pending_form_id="form-1", pending_tool_call_id="call-1", ) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py b/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py index f65770665e4..d8397fe3a55 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py @@ -9,13 +9,18 @@ from __future__ import annotations from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock import pytest from core.app.apps.agent_app.app_generator import AgentAppGenerator, AgentAppGeneratorError, AgentAppNotPublishedError from core.app.entities.app_invoke_entities import InvokeFrom -from models.agent import AgentConfigDraft, AgentConfigDraftType, AgentScope, AgentSource +from models.agent import AgentConfigDraft, AgentConfigDraftType, AgentConfigVersionKind, AgentScope, AgentSource from models.agent_config_entities import AgentSoulConfig +from services.agent.workspace_service import ( + AgentWorkspaceBindingGenerationMismatchError, + AgentWorkspaceService, +) _SOUL_DICT = { "model": { @@ -34,8 +39,10 @@ class _FakeScalarSession: self._values = list(values) self.added: list[Any] = [] self.flush_count = 0 + self.scalar_statements: list[Any] = [] - def scalar(self, _stmt: Any) -> Any: + def scalar(self, stmt: Any) -> Any: + self.scalar_statements.append(stmt) return self._values.pop(0) if self._values else None def add(self, value: Any) -> None: @@ -46,7 +53,7 @@ class _FakeScalarSession: def _snapshot() -> SimpleNamespace: - return SimpleNamespace(id="snap-1", config_snapshot_dict=_SOUL_DICT) + return SimpleNamespace(id="snap-1", home_snapshot_id="home-1", config_snapshot_dict=_SOUL_DICT) class TestResolveAgentById: @@ -102,6 +109,7 @@ class TestResolveDebugDraft: tenant_id="t1", agent=agent, draft_type=None, + draft_id=None, account_id=None, session=session, ) @@ -127,10 +135,12 @@ class TestResolveDebugDraft: account_id=None, draft_owner_key="", base_snapshot_id="snap-1", + home_snapshot_id="home-1", config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "old"}}), ) active_snapshot = SimpleNamespace( id="snap-2", + home_snapshot_id="home-2", config_snapshot_dict={"prompt": {"system_prompt": "new"}}, ) session = _FakeScalarSession([draft, active_snapshot]) @@ -146,6 +156,7 @@ class TestResolveDebugDraft: assert resolved is draft assert resolved.id == "draft-1" assert resolved.base_snapshot_id == "snap-2" + assert resolved.home_snapshot_id == "home-2" assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "new" assert session.flush_count == 1 @@ -165,6 +176,7 @@ class TestResolveDebugDraft: account_id="account-1", draft_owner_key="account-1", base_snapshot_id="snap-1", + home_snapshot_id="home-build", config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "build edit"}}), ) session = _FakeScalarSession([draft]) @@ -182,6 +194,48 @@ class TestResolveDebugDraft: assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "build edit" assert session.flush_count == 0 + def test_build_draft_uses_exact_draft_id(self): + agent = SimpleNamespace( + id="agent-1", + scope=AgentScope.WORKFLOW_ONLY, + active_config_snapshot_id="snap-2", + created_by="creator-1", + updated_by="updater-1", + ) + draft = AgentConfigDraft( + id="exact-build-draft", + tenant_id="t1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + base_snapshot_id="snap-1", + home_snapshot_id="home-build", + config_snapshot=AgentSoulConfig(), + ) + session = _FakeScalarSession([draft]) + statements: list[Any] = [] + scalar = session.scalar + + def capture_scalar(statement: Any) -> Any: + statements.append(statement) + return scalar(statement) + + session.scalar = capture_scalar # type: ignore[method-assign] + + resolved = AgentAppGenerator._resolve_debug_draft( + tenant_id="t1", + agent=agent, + draft_type=AgentConfigDraftType.DEBUG_BUILD.value, + draft_id="exact-build-draft", + account_id="account-1", + session=session, + ) + + assert resolved is draft + assert "agent_config_drafts.id =" in str(statements[0]) + assert "exact-build-draft" in statements[0].compile().params.values() + class TestResolveAgent: def test_success_chains_to_resolve_by_id(self): @@ -235,6 +289,126 @@ class TestResolveAgent: assert config_version_kind == "snapshot" assert soul.prompt.system_prompt == "You are Iris." + def test_existing_conversation_resolves_binding_snapshot_instead_of_latest_active_snapshot( + self, monkeypatch: pytest.MonkeyPatch + ): + bound_agent = SimpleNamespace( + id="agent-1", + source=AgentSource.AGENT_APP, + active_config_snapshot_id="snap-2", + active_config_is_published=True, + ) + inner_agent = SimpleNamespace(id="agent-1") + pinned_snapshot = SimpleNamespace( + id="snap-1", + home_snapshot_id="home-1", + config_snapshot_dict=_SOUL_DICT, + ) + conversation = SimpleNamespace(id="conversation-1", agent_workspace_binding_id="binding-1") + binding = SimpleNamespace( + id="binding-1", + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="snap-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + get_active_binding = MagicMock(return_value=binding) + validate_generation = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) + monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) + session = _FakeScalarSession([bound_agent, inner_agent, pinned_snapshot]) + app_model = SimpleNamespace(id="app-1", tenant_id="t1") + + _, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=SimpleNamespace(id="user-1"), + session=session, + conversation=conversation, + ) # type: ignore[arg-type] + + assert config_id == "snap-1" + assert config_version_kind == "snapshot" + assert soul.prompt.system_prompt == "You are Iris." + assert get_active_binding.call_args.kwargs["binding_id"] == "binding-1" + validate_generation.assert_called_once_with( + binding, + base_home_snapshot_id="home-1", + agent_config_version_id="snap-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + + def test_existing_conversation_rejects_unavailable_binding(self, monkeypatch: pytest.MonkeyPatch): + bound_agent = SimpleNamespace( + id="agent-1", + source=AgentSource.AGENT_APP, + active_config_snapshot_id="snap-active", + active_config_is_published=True, + ) + conversation = SimpleNamespace(id="conversation-1", agent_workspace_binding_id="binding-missing") + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None)) + + with pytest.raises(AgentAppGeneratorError, match="Conversation participant Binding is unavailable"): + AgentAppGenerator()._resolve_agent( + SimpleNamespace(id="app-1", tenant_id="t1"), + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=SimpleNamespace(id="user-1"), + session=_FakeScalarSession([bound_agent]), + conversation=conversation, + ) # type: ignore[arg-type] + + @pytest.mark.parametrize( + ("binding_home_id", "binding_version_kind", "snapshot_home_id"), + [ + ("home-binding", AgentConfigVersionKind.SNAPSHOT, "home-other"), + ("home-pinned", AgentConfigVersionKind.DRAFT, "home-pinned"), + ], + ) + def test_existing_conversation_generation_mismatch_does_not_fallback_to_active_snapshot( + self, + monkeypatch: pytest.MonkeyPatch, + binding_home_id: str, + binding_version_kind: AgentConfigVersionKind, + snapshot_home_id: str, + ): + bound_agent = SimpleNamespace( + id="agent-1", + source=AgentSource.AGENT_APP, + active_config_snapshot_id="snap-active", + active_config_is_published=True, + ) + conversation = SimpleNamespace(id="conversation-1", agent_workspace_binding_id="binding-1") + binding = SimpleNamespace( + id="binding-1", + agent_id="agent-1", + base_home_snapshot_id=binding_home_id, + agent_config_version_id="snap-pinned", + agent_config_version_kind=binding_version_kind, + ) + pinned_snapshot = SimpleNamespace( + id="snap-pinned", + home_snapshot_id=snapshot_home_id, + config_snapshot_dict=_SOUL_DICT, + ) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) + session = _FakeScalarSession([bound_agent, SimpleNamespace(id="agent-1"), pinned_snapshot]) + + with pytest.raises(AgentWorkspaceBindingGenerationMismatchError): + AgentAppGenerator()._resolve_agent( + SimpleNamespace(id="app-1", tenant_id="t1"), + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=SimpleNamespace(id="user-1"), + session=session, + conversation=conversation, + ) # type: ignore[arg-type] + + snapshot_query_params = session.scalar_statements[-1].compile().params.values() + assert "snap-pinned" in snapshot_query_params + assert "snap-active" not in snapshot_query_params + def test_unpublished_imported_agent_is_not_available_to_public_runtime(self): bound_agent = SimpleNamespace( id="agent-1", diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py index b3ea05b8a87..a027fd31699 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py @@ -44,6 +44,7 @@ class TestBuildForAgentApp: AgentBackendAgentAppRunInput( model=AgentBackendModelConfig(plugin_id="langgenius/openai", model_provider="openai", model="gpt-test"), execution_context=_exec_ctx(), + backend_binding_ref="binding-ref-1", user_prompt="hello", agent_soul_prompt="You are Iris.", ) @@ -65,6 +66,7 @@ class TestBuildForAgentApp: AgentBackendAgentAppRunInput( model=AgentBackendModelConfig(plugin_id="p/q", model_provider="openai", model="m"), execution_context=_exec_ctx(), + backend_binding_ref="binding-ref-1", user_prompt=" ", ) @@ -73,6 +75,7 @@ class TestBuildForAgentApp: AgentBackendAgentAppRunInput( model=AgentBackendModelConfig(plugin_id="langgenius/openai", model_provider="openai", model="gpt-test"), execution_context=_exec_ctx(), + backend_binding_ref="binding-ref-1", user_prompt="hi", ) ) @@ -142,7 +145,6 @@ def _ctx( *, query: str = "hello", agent_config_version_kind: str = "snapshot", - suspend_on_exit: bool = True, ) -> AgentAppRuntimeBuildContext: dify_context = SimpleNamespace( tenant_id="tenant-1", @@ -159,8 +161,9 @@ def _ctx( conversation_id="conv-1", user_query=query, idempotency_key="msg-1", + binding_id="binding-1", + backend_binding_ref="binding-ref-1", agent_config_version_kind=agent_config_version_kind, # type: ignore[arg-type] - suspend_on_exit=suspend_on_exit, ) @@ -191,6 +194,7 @@ class TestAgentAppRuntimeRequestBuilder: "agent_soul_prompt", "agent_app_user_prompt", "execution_context", + "runtime", DIFY_SHELL_LAYER_ID, DIFY_CONFIG_LAYER_ID, "history", @@ -243,16 +247,6 @@ class TestAgentAppRuntimeRequestBuilder: assert execution_context.config.agent_config_version_kind == "draft" assert config_layer.config.config_version.kind == "draft" - def test_build_uses_delete_on_exit_when_requested(self): - builder = AgentAppRuntimeRequestBuilder( - credentials_provider=_FakeCredentialsProvider(), - dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] - ) - - result = builder.build(_ctx(_soul_with_model(), suspend_on_exit=False)) - - assert result.request.on_exit.default.value == "delete" - def test_build_includes_plugin_tools_layer_returned_by_injected_builder_for_draft(self): soul = _soul_with_model() soul.tools.dify_tools = [ @@ -410,7 +404,7 @@ class TestAgentAppRuntimeRequestBuilder: ] assert shell_config["cli_tools"][0]["install_commands"] == ["apt-get install -y ripgrep"] assert shell_config["env"][0] == {"name": "PROJECT_NAME", "value": "demo"} - assert shell_config["sandbox"] == {"provider": "independent", "config": {"cpu": 2}} + assert "sandbox" not in shell_config assert result.metadata["agent_tools"] == { "dify_tool_count": 0, "dify_tool_names": [], @@ -461,7 +455,7 @@ class TestAgentAppConfigLayer: assert config.config.mentioned_file_names == [] # shell enters first; config uses that shell to materialize mentioned targets. names = [layer.name for layer in result.request.composition.layers] - assert names.index(DIFY_SHELL_LAYER_ID) == names.index("execution_context") + 1 + assert names.index(DIFY_SHELL_LAYER_ID) == names.index("execution_context") + 2 assert names.index(DIFY_CONFIG_LAYER_ID) == names.index(DIFY_SHELL_LAYER_ID) + 1 def test_config_layer_present_when_agent_soul_has_no_config_assets(self, monkeypatch: pytest.MonkeyPatch): @@ -484,7 +478,10 @@ class TestAgentAppConfigLayer: "mentioned_skill_names": [], "mentioned_file_names": [], } - assert layers[DIFY_SHELL_LAYER_ID].deps == {"execution_context": "execution_context"} + assert layers[DIFY_SHELL_LAYER_ID].deps == { + "execution_context": "execution_context", + "runtime": "runtime", + } assert layers[DIFY_SHELL_LAYER_ID].config.agent_stub_drive_ref is None def test_config_layer_for_build_draft_marks_config_writable(self): diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py b/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py index fc9bf1b48e6..a604637eb80 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py @@ -1,348 +1,123 @@ -"""Unit tests for the conversation-keyed Agent App session store. - -Exercises the real ORM round-trip against the project's in-memory SQLite engine -(per-test create/drop of the unified ``agent_runtime_sessions`` table), so the -conversation owner path is verified without Postgres. -""" - -from __future__ import annotations - -from collections.abc import Generator +from types import SimpleNamespace +from unittest.mock import MagicMock import pytest from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from dify_agent.protocol import RuntimeLayerSpec -from sqlalchemy import delete -from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore, AgentAppSessionScope -from core.db.session_factory import session_factory -from models.agent import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus +from core.app.apps.agent_app.session_store import AgentAppSessionScope, AgentAppWorkspaceStore +from models.agent import ( + AgentConfigVersionKind, + AgentWorkspaceOwnerType, +) +from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService def _scope( - conversation_id: str = "conv-1", agent_id: str = "agent-1", agent_config_snapshot_id: str | None = "snap-1" + *, + kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT, + build_draft_id: str | None = None, ) -> AgentAppSessionScope: return AgentAppSessionScope( tenant_id="tenant-1", app_id="app-1", - conversation_id=conversation_id, - agent_id=agent_id, - agent_config_snapshot_id=agent_config_snapshot_id, + conversation_id="conversation-1", + agent_id="agent-1", + agent_config_snapshot_id="config-1", + home_snapshot_id="home-1", + agent_config_version_kind=kind, + build_draft_id=build_draft_id, ) -def _snapshot(messages: int = 1) -> CompositorSessionSnapshot: - return CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="history", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={"messages": [{"role": "user", "content": f"m{i}"} for i in range(messages)]}, - ) - ] +def _binding() -> SimpleNamespace: + return SimpleNamespace( + id="binding-1", + workspace_id="workspace-1", + backend_binding_ref="backend-binding-1", + agent_id="agent-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + base_home_snapshot_id="home-1", + session_snapshot=None, + pending_form_id=None, + pending_tool_call_id=None, ) -def _runtime_layer_specs() -> list[RuntimeLayerSpec]: - return [ - RuntimeLayerSpec(name="execution_context", type="dify.execution_context", config={"tenant_id": "tenant-1"}), - RuntimeLayerSpec(name="history", type="pydantic_ai.history"), - ] +def test_scope_selects_conversation_or_build_draft_workspace_owner() -> None: + assert _scope().workspace_owner.owner_type is AgentWorkspaceOwnerType.CONVERSATION + build_owner = _scope( + kind=AgentConfigVersionKind.BUILD_DRAFT, + build_draft_id="build-draft-1", + ).workspace_owner + assert build_owner.owner_type is AgentWorkspaceOwnerType.BUILD_DRAFT + assert build_owner.owner_id == "build-draft-1" -@pytest.fixture(autouse=True) -def _create_table() -> Generator[None, None, None]: - engine = session_factory.get_session_maker().kw["bind"] - AgentRuntimeSession.__table__.create(bind=engine, checkfirst=True) - yield - with session_factory.create_session() as session: - session.execute(delete(AgentRuntimeSession)) - session.commit() - AgentRuntimeSession.__table__.drop(bind=engine, checkfirst=True) +def test_load_or_create_persists_new_binding_on_caller(monkeypatch) -> None: + caller = SimpleNamespace(agent_workspace_binding_id=None) + context = MagicMock() + session = context.__enter__.return_value + create = MagicMock(return_value=_binding()) + store = AgentAppWorkspaceStore() + monkeypatch.setattr("core.app.apps.agent_app.session_store.session_factory.create_session", lambda: context) + monkeypatch.setattr(store, "_load_caller", MagicMock(return_value=caller)) + monkeypatch.setattr(AgentWorkspaceService, "create_binding", create) + + stored = store.load_or_create(_scope()) + + assert stored.binding_id == "binding-1" + assert stored.workspace_id == "workspace-1" + assert stored.backend_binding_ref == "backend-binding-1" + assert caller.agent_workspace_binding_id == "binding-1" + assert create.call_args.kwargs["session"] is session + session.commit.assert_called_once_with() -def test_load_returns_none_when_no_row(): - assert AgentAppRuntimeSessionStore().load_active_snapshot(_scope()) is None +def test_load_or_create_uses_exact_caller_binding(monkeypatch) -> None: + caller = SimpleNamespace(agent_workspace_binding_id="binding-1") + context = MagicMock() + context.__enter__.return_value = MagicMock() + get_binding = MagicMock(return_value=_binding()) + create = MagicMock() + store = AgentAppWorkspaceStore() + monkeypatch.setattr("core.app.apps.agent_app.session_store.session_factory.create_session", lambda: context) + monkeypatch.setattr(store, "_load_caller", MagicMock(return_value=caller)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding) + monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", MagicMock()) + monkeypatch.setattr(AgentWorkspaceService, "create_binding", create) + + stored = store.load_or_create(_scope()) + + assert stored.binding_id == "binding-1" + assert get_binding.call_args.kwargs["binding_id"] == "binding-1" + create.assert_not_called() -def test_save_creates_conversation_owned_row_and_round_trips(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(messages=2), - runtime_layer_specs=_runtime_layer_specs(), - ) +def test_normal_conversation_pointer_does_not_create_replacement_binding(monkeypatch) -> None: + caller = SimpleNamespace(agent_workspace_binding_id="unavailable-binding") + context = MagicMock() + get_binding = MagicMock(return_value=None) + create = MagicMock() + store = AgentAppWorkspaceStore() + monkeypatch.setattr("core.app.apps.agent_app.session_store.session_factory.create_session", lambda: context) + monkeypatch.setattr(store, "_load_caller", MagicMock(return_value=caller)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding) + monkeypatch.setattr(AgentWorkspaceService, "create_binding", create) - loaded = store.load_active_snapshot(_scope()) - assert loaded is not None - assert loaded.layers[0].runtime_state["messages"] == [ - {"role": "user", "content": "m0"}, - {"role": "user", "content": "m1"}, - ] - with session_factory.create_session() as session: - row = session.query(AgentRuntimeSession).one() - assert row.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION - assert row.conversation_id == "conv-1" - assert row.agent_config_snapshot_id == "snap-1" - assert row.workflow_run_id is None # conversation owner leaves workflow cols NULL - assert row.backend_run_id == "run-1" - assert "execution_context" in row.composition_layer_specs - assert "history" in row.composition_layer_specs + with pytest.raises(AgentWorkspaceNotFoundError, match="Caller participant Binding is unavailable"): + store.load_or_create(_scope()) + + assert get_binding.call_args.kwargs["binding_id"] == "unavailable-binding" + create.assert_not_called() -def test_save_is_noop_when_snapshot_missing(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-x", - snapshot=None, - runtime_layer_specs=_runtime_layer_specs(), - ) - with session_factory.create_session() as session: - assert session.query(AgentRuntimeSession).count() == 0 +def test_save_snapshot_targets_binding(monkeypatch) -> None: + save = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "save_binding_session_snapshot", save) + snapshot = CompositorSessionSnapshot(layers=[]) + AgentAppWorkspaceStore().save_active_snapshot(scope=_scope(), binding_id="binding-1", snapshot=snapshot) -def test_second_turn_updates_same_conversation_row(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(messages=1), - runtime_layer_specs=_runtime_layer_specs(), - ) - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-2", - snapshot=_snapshot(messages=3), - runtime_layer_specs=_runtime_layer_specs(), - ) - with session_factory.create_session() as session: - rows = session.query(AgentRuntimeSession).all() - assert len(rows) == 1 - assert rows[0].backend_run_id == "run-2" - - -def test_debug_scope_with_null_snapshot_id_updates_same_conversation_row(): - store = AgentAppRuntimeSessionStore() - scope = _scope(agent_config_snapshot_id=None) - store.save_active_snapshot( - scope=scope, - backend_run_id="run-1", - snapshot=_snapshot(messages=1), - runtime_layer_specs=_runtime_layer_specs(), - ) - store.save_active_snapshot( - scope=scope, - backend_run_id="run-2", - snapshot=_snapshot(messages=3), - runtime_layer_specs=_runtime_layer_specs(), - ) - - loaded = store.load_active_snapshot(scope) - - assert loaded is not None - assert loaded.layers[0].runtime_state["messages"] == [ - {"role": "user", "content": "m0"}, - {"role": "user", "content": "m1"}, - {"role": "user", "content": "m2"}, - ] - with session_factory.create_session() as session: - row = session.query(AgentRuntimeSession).one() - assert row.agent_config_snapshot_id is None - assert row.backend_run_id == "run-2" - - -def test_mark_cleaned_then_load_returns_none_and_save_resurrects(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(), - runtime_layer_specs=_runtime_layer_specs(), - ) - store.mark_cleaned(scope=_scope(), backend_run_id="cleanup-1") - assert store.load_active_snapshot(_scope()) is None - # Re-entry revives the row. - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-2", - snapshot=_snapshot(messages=2), - runtime_layer_specs=_runtime_layer_specs(), - ) - with session_factory.create_session() as session: - row = session.query(AgentRuntimeSession).one() - assert row.status == AgentRuntimeSessionStatus.ACTIVE - assert row.cleaned_at is None - assert row.backend_run_id == "run-2" - - -def test_distinct_conversations_do_not_collide(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(conversation_id="conv-A"), - backend_run_id="a", - snapshot=_snapshot(), - runtime_layer_specs=_runtime_layer_specs(), - ) - store.save_active_snapshot( - scope=_scope(conversation_id="conv-B"), - backend_run_id="b", - snapshot=_snapshot(), - runtime_layer_specs=_runtime_layer_specs(), - ) - assert store.load_active_snapshot(_scope(conversation_id="conv-A")) is not None - assert store.load_active_snapshot(_scope(conversation_id="conv-B")) is not None - with session_factory.create_session() as session: - assert session.query(AgentRuntimeSession).count() == 2 - - -def test_distinct_agent_config_snapshots_keep_only_latest_active_session(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(agent_config_snapshot_id="snap-1"), - backend_run_id="a", - snapshot=_snapshot(), - runtime_layer_specs=_runtime_layer_specs(), - ) - store.save_active_snapshot( - scope=_scope(agent_config_snapshot_id="snap-2"), - backend_run_id="b", - snapshot=_snapshot(messages=2), - runtime_layer_specs=_runtime_layer_specs(), - ) - - assert store.load_active_snapshot(_scope(agent_config_snapshot_id="snap-1")) is None - assert store.load_active_snapshot(_scope(agent_config_snapshot_id="snap-2")) is not None - with session_factory.create_session() as session: - rows = session.query(AgentRuntimeSession).order_by(AgentRuntimeSession.backend_run_id).all() - assert len(rows) == 2 - assert [row.agent_config_snapshot_id for row in rows] == ["snap-1", "snap-2"] - assert [row.status for row in rows] == [AgentRuntimeSessionStatus.CLEANED, AgentRuntimeSessionStatus.ACTIVE] - - -def test_load_active_session_for_conversation_resolves_without_agent_or_config_scope(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(messages=2), - runtime_layer_specs=_runtime_layer_specs(), - ) - - loaded = store.load_active_session_for_conversation(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1") - assert loaded is not None - assert loaded.session_snapshot.layers[0].runtime_state["messages"] == [ - {"role": "user", "content": "m0"}, - {"role": "user", "content": "m1"}, - ] - assert [spec.name for spec in loaded.runtime_layer_specs] == ["execution_context", "history"] - - -def test_load_active_session_for_conversation_uses_latest_active_snapshot_after_config_change(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(agent_config_snapshot_id="snap-1"), - backend_run_id="a", - snapshot=_snapshot(), - runtime_layer_specs=_runtime_layer_specs(), - ) - store.save_active_snapshot( - scope=_scope(agent_config_snapshot_id="snap-2"), - backend_run_id="b", - snapshot=_snapshot(messages=3), - runtime_layer_specs=_runtime_layer_specs(), - ) - - loaded = store.load_active_session_for_conversation(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1") - - assert loaded is not None - assert loaded.session_snapshot.layers[0].runtime_state["messages"] == [ - {"role": "user", "content": "m0"}, - {"role": "user", "content": "m1"}, - {"role": "user", "content": "m2"}, - ] - - -def test_load_active_session_for_conversation_returns_none_when_cleaned_or_absent(): - store = AgentAppRuntimeSessionStore() - assert ( - store.load_active_session_for_conversation(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1") - is None - ) - - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(), - runtime_layer_specs=_runtime_layer_specs(), - ) - store.mark_cleaned(scope=_scope(), backend_run_id="cleanup-1") - assert ( - store.load_active_session_for_conversation(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1") - is None - ) - - -def test_load_active_session_for_conversation_isolates_other_conversations(): - store = AgentAppRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(conversation_id="conv-A"), - backend_run_id="a", - snapshot=_snapshot(), - runtime_layer_specs=_runtime_layer_specs(), - ) - - assert ( - store.load_active_session_for_conversation(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-B") - is None - ) - assert ( - store.load_active_session_for_conversation(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-A") - is not None - ) - - -def test_list_active_sessions_for_conversation_returns_all_active_rows(): - store = AgentAppRuntimeSessionStore() - with session_factory.create_session() as session: - session.add( - AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-1", - agent_config_snapshot_id="snap-1", - conversation_id="conv-1", - backend_run_id="run-1", - session_snapshot=_snapshot(messages=1).model_dump_json(), - composition_layer_specs='[{"name":"execution_context","type":"dify.execution_context","deps":{},"metadata":{},"config":{"tenant_id":"tenant-1"}},{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', - status=AgentRuntimeSessionStatus.ACTIVE, - ) - ) - session.add( - AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-2", - agent_config_snapshot_id="snap-2", - conversation_id="conv-1", - backend_run_id="run-2", - session_snapshot=_snapshot(messages=2).model_dump_json(), - composition_layer_specs='[{"name":"execution_context","type":"dify.execution_context","deps":{},"metadata":{},"config":{"tenant_id":"tenant-1"}},{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', - status=AgentRuntimeSessionStatus.ACTIVE, - ) - ) - session.commit() - - loaded = store.list_active_sessions_for_conversation( - tenant_id="tenant-1", - app_id="app-1", - conversation_id="conv-1", - ) - - assert [session.scope.agent_id for session in loaded] == ["agent-2", "agent-1"] - assert all(session.scope.conversation_id == "conv-1" for session in loaded) + assert save.call_args.kwargs["binding_id"] == "binding-1" + assert save.call_args.kwargs["session_snapshot"] == snapshot.model_dump_json() diff --git a/api/tests/unit_tests/core/app/workflow/layers/test_persistence.py b/api/tests/unit_tests/core/app/workflow/layers/test_persistence.py index 5c50cb78dae..b7a8754d92c 100644 --- a/api/tests/unit_tests/core/app/workflow/layers/test_persistence.py +++ b/api/tests/unit_tests/core/app/workflow/layers/test_persistence.py @@ -34,6 +34,7 @@ def test_update_node_execution_prefers_event_finished_at(monkeypatch: pytest.Mon node_execution = Mock() node_execution.id = "node-exec-1" node_execution.created_at = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC).replace(tzinfo=None) + node_execution.process_data = None node_execution.update_from_mapping = Mock() layer._node_snapshots[node_execution.id] = _NodeRuntimeSnapshot( @@ -66,6 +67,7 @@ def test_update_node_execution_projects_start_outputs() -> None: node_execution.id = "node-exec-2" node_execution.node_type = BuiltinNodeTypes.START node_execution.created_at = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC).replace(tzinfo=None) + node_execution.process_data = None node_execution.update_from_mapping = Mock() layer._node_snapshots[node_execution.id] = _NodeRuntimeSnapshot( diff --git a/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py b/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py index 31b9bb28940..52d0db414c2 100644 --- a/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py +++ b/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py @@ -39,11 +39,15 @@ from graphon.runtime import GraphRuntimeState, ReadOnlyGraphRuntimeStateWrapper, class _RepoRecorder: def __init__(self) -> None: self.saved: list[object] = [] + self.synchronously_saved: list[object] = [] self.saved_exec_data: list[object] = [] def save(self, entity): self.saved.append(entity) + def save_synchronously(self, entity): + self.synchronously_saved.append(entity) + def save_execution_data(self, entity): self.saved_exec_data.append(entity) @@ -331,6 +335,24 @@ class TestWorkflowPersistenceLayer: layer._handle_node_retry(retry_event) assert node_repo.saved_exec_data + def test_agent_v2_caller_row_is_saved_synchronously_before_node_run(self): + layer, _, node_repo, _ = _make_layer() + layer._handle_graph_run_started() + + layer._handle_node_started( + NodeRunStartedEvent( + id="agent-exec", + node_id="agent-node", + node_type=BuiltinNodeTypes.AGENT, + node_version="2", + node_title="Agent", + start_at=_naive_utc_now(), + ) + ) + + assert [execution.id for execution in node_repo.synchronously_saved] == ["agent-exec"] + assert node_repo.saved == [] + def test_retry_history_is_preserved_after_node_succeeds(self): layer, _, node_repo, _ = _make_layer() layer._handle_graph_run_started() @@ -501,7 +523,12 @@ class TestWorkflowPersistenceLayer: domain_execution = layer._node_execution_cache["exec"] domain_execution.inputs = {"old": True} - result = NodeRunResult(inputs={"new": True}, outputs={"out": 1}, process_data={"p": 1}, metadata={}) + result = NodeRunResult( + inputs={"new": True}, + outputs={"out": 1}, + process_data={"p": 1, "workflow_agent_binding_id": "workflow-binding-1"}, + metadata={}, + ) pause_event = NodeRunPauseRequestedEvent( id="exec", node_id="node", @@ -513,6 +540,38 @@ class TestWorkflowPersistenceLayer: assert domain_execution.status == WorkflowNodeExecutionStatus.PAUSED assert domain_execution.inputs == {"old": True} + assert domain_execution.process_data == {"workflow_agent_binding_id": "workflow-binding-1"} + + def test_handle_node_retry_preserves_workflow_agent_binding_identity(self): + layer, _, _, _ = _make_layer() + layer._handle_graph_run_started() + started_at = _naive_utc_now() + layer._handle_node_started( + NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.AGENT, + node_title="Agent", + start_at=started_at, + ) + ) + + layer._handle_node_retry( + NodeRunRetryEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.AGENT, + node_title="Agent", + start_at=started_at, + error="retry", + retry_index=1, + node_run_result=NodeRunResult( + process_data={"workflow_agent_binding_id": "workflow-binding-1"}, + ), + ) + ) + + assert layer._node_execution_cache["exec"].process_data["workflow_agent_binding_id"] == "workflow-binding-1" def test_get_node_execution_raises_for_missing(self): layer, _, _, _ = _make_layer() diff --git a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py index 018626c450c..b502144b0cf 100644 --- a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py @@ -186,6 +186,23 @@ class TestCeleryWorkflowNodeExecutionRepository: in repo._workflow_execution_mapping[sample_workflow_node_execution.workflow_execution_id] ) + def test_save_synchronously_uses_sql_repository_without_queueing( + self, mock_session_factory, mock_account, sample_workflow_node_execution + ): + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=mock_account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + repo._sql_repository.save_synchronously = Mock() + + repo.save_synchronously(sample_workflow_node_execution) + + repo._sql_repository.save_synchronously.assert_called_once_with(sample_workflow_node_execution) + assert repo._execution_cache[sample_workflow_node_execution.id] is sample_workflow_node_execution + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") def test_save_handles_celery_failure( self, mock_task, mock_session_factory, mock_account, sample_workflow_node_execution diff --git a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py index 3289e2b50e8..bf750ab68d8 100644 --- a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py @@ -238,7 +238,11 @@ def test_to_db_model_requires_constructor_context(monkeypatch: pytest.MonkeyPatc app_id=None, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) - execution = _execution(inputs={"b": 1, "a": 2}, metadata={WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 1}) + execution = _execution( + inputs={"b": 1, "a": 2}, + process_data={"agent_workspace_binding_id": "participant-1"}, + metadata={WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 1}, + ) # Happy path: deterministic json dump should be sorted db_model = repo._to_db_model(execution) @@ -247,6 +251,7 @@ def test_to_db_model_requires_constructor_context(monkeypatch: pytest.MonkeyPatc assert db_model.created_by_role.value == "account" assert json.loads(db_model.inputs or "{}") == {"a": 2, "b": 1} assert json.loads(db_model.execution_metadata or "{}")["total_tokens"] == 1 + assert db_model.agent_workspace_binding_id is None repo._triggered_from = None with pytest.raises(ValueError, match="triggered_from is required"): @@ -330,7 +335,7 @@ def test_persist_to_database_updates_existing_and_inserts_new(monkeypatch: pytes db_model.foo = "bar" # type: ignore[attr-defined] db_model.__dict__["_private"] = "x" - existing = SimpleNamespace() + existing = SimpleNamespace(process_data=None, process_data_dict=None) session.get.return_value = existing repo._persist_to_database(db_model) assert existing.foo == "bar" @@ -523,7 +528,9 @@ def test_save_execution_data_handles_existing_db_model_and_truncation(monkeypatc offload_data=[WorkflowNodeExecutionOffload(type_=ExecutionOffLoadType.INPUTS)], inputs=None, outputs=None, - process_data=None, + process_data='{"workflow_agent_binding_id": "workflow-binding-1"}', + process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"}, + agent_workspace_binding_id="authoritative-participant", ) session.merge = Mock() session.flush = Mock() @@ -538,7 +545,11 @@ def test_save_execution_data_handles_existing_db_model_and_truncation(monkeypatc triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) - execution = _execution(inputs={"a": 1}, outputs={"b": 2}, process_data={"c": 3}) + execution = _execution( + inputs={"a": 1}, + outputs={"b": 2}, + process_data={"c": 3}, + ) trunc_result = SimpleNamespace( truncated_value={"trunc": True}, @@ -554,7 +565,11 @@ def test_save_execution_data_handles_existing_db_model_and_truncation(monkeypatc db_model = session.merge.call_args.args[0] assert json.loads(db_model.inputs) == {"trunc": True} assert json.loads(db_model.outputs) == {"b": 2} - assert json.loads(db_model.process_data) == {"c": 3} + assert json.loads(db_model.process_data) == { + "c": 3, + "workflow_agent_binding_id": "workflow-binding-1", + } + assert db_model.agent_workspace_binding_id == "authoritative-participant" assert any(off.type_ == ExecutionOffLoadType.INPUTS for off in db_model.offload_data) assert execution.get_truncated_inputs() == {"trunc": True} @@ -570,6 +585,7 @@ def test_save_execution_data_truncates_outputs_and_process_data(monkeypatch: pyt inputs=None, outputs=None, process_data=None, + process_data_dict=None, ) session = MagicMock() session.execute.return_value.scalars.return_value.first.return_value = existing @@ -633,7 +649,14 @@ def test_save_execution_data_handles_missing_db_model(monkeypatch: pytest.Monkey ) execution = _execution(inputs={"a": 1}) - fake_db_model = SimpleNamespace(id=execution.id, offload_data=[], inputs=None, outputs=None, process_data=None) + fake_db_model = SimpleNamespace( + id=execution.id, + offload_data=[], + inputs=None, + outputs=None, + process_data=None, + process_data_dict=None, + ) monkeypatch.setattr(repo, "_to_db_model", lambda *_: fake_db_model) monkeypatch.setattr(repo, "_truncate_and_upload", lambda *_args, **_kwargs: None) monkeypatch.setattr(repo, "_json_encode", lambda values: json.dumps(values)) diff --git a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py index e70745f70b3..f22102bd427 100644 --- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py +++ b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py @@ -99,6 +99,7 @@ class TestWorkflowNodeExecutionConflictHandling: # Mock existing record mock_existing = MagicMock() + mock_existing.process_data_dict = None mock_session.get.return_value = mock_existing mock_session.commit.return_value = None diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py index d142591c3fb..4dd7dc9fd80 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime from types import SimpleNamespace from typing import cast from unittest.mock import MagicMock, patch +from uuid import UUID from agenton.compositor import CompositorSessionSnapshot from dify_agent.layers.ask_human import AskHumanToolResult @@ -24,7 +25,6 @@ from clients.agent_backend import ( AgentBackendStreamInternalEvent, FakeAgentBackendRunClient, FakeAgentBackendScenario, - RuntimeLayerSpec, ) from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext, InvokeFrom, UserFrom from core.workflow.file_reference import build_file_reference @@ -33,18 +33,22 @@ from core.workflow.nodes.agent_v2.ask_human_resume import AskHumanResumeOutcome from core.workflow.nodes.agent_v2.binding_resolver import WorkflowAgentBindingBundle, WorkflowAgentBindingResolver from core.workflow.nodes.agent_v2.entities import DifyAgentNodeData from core.workflow.nodes.agent_v2.output_adapter import WorkflowAgentOutputAdapter -from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeRequestBuilder +from core.workflow.nodes.agent_v2.runtime_request_builder import ( + WorkflowAgentRuntimeBuildContext, + WorkflowAgentRuntimeRequestBuilder, +) from core.workflow.nodes.agent_v2.session_store import ( StoredWorkflowAgentSession, - WorkflowAgentRuntimeSessionStore, WorkflowAgentSessionScope, + WorkflowAgentWorkspaceStore, ) from core.workflow.nodes.human_input.pause_reason import HumanInputRequired from graphon.entities import GraphInitParams from graphon.entities.pause_reason import HitlRequired from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus from graphon.file import File, FileTransferMethod, FileType -from graphon.node_events import PauseRequestedEvent, StreamCompletedEvent +from graphon.graph_events import NodeRunPauseRequestedEvent +from graphon.node_events import StreamCompletedEvent from graphon.runtime import GraphRuntimeState from graphon.variables.segments import ArrayFileSegment, FileSegment, StringSegment from models.agent import Agent, AgentConfigSnapshot, WorkflowAgentNodeBinding @@ -54,6 +58,7 @@ from models.agent_config_entities import ( DeclaredOutputType, WorkflowNodeJobConfig, ) +from services.agent.workspace_service import AgentWorkspaceNotFoundError class FakeCredentialsProvider: @@ -95,12 +100,14 @@ class FakeVariablePool: class FakeBindingResolver(WorkflowAgentBindingResolver): def __init__(self): + self.calls: list[dict[str, object]] = [] self.agent = Agent(id="agent-1", tenant_id="tenant-1", name="Agent") self.snapshot = AgentConfigSnapshot( id="snapshot-1", tenant_id="tenant-1", agent_id="agent-1", version=1, + home_snapshot_id="home-1", config_snapshot=AgentSoulConfig( prompt={"system_prompt": "You are careful."}, model=AgentSoulModelConfig( @@ -127,13 +134,29 @@ class FakeBindingResolver(WorkflowAgentBindingResolver): ), ) - def resolve(self, **_kwargs): + def resolve(self, **kwargs): + self.calls.append(kwargs) + snapshot_id = kwargs.get("snapshot_id") + if isinstance(snapshot_id, str): + self.snapshot.id = snapshot_id return WorkflowAgentBindingBundle(binding=self.binding, agent=self.agent, snapshot=self.snapshot) class FakeSessionStore: - def __init__(self, snapshot: CompositorSessionSnapshot | None = None) -> None: + def __init__( + self, + snapshot: CompositorSessionSnapshot | None = None, + *, + binding_id: str = "binding-1", + workspace_id: str = "workspace-1", + backend_binding_ref: str = "backend-binding-1", + ) -> None: self.loaded_snapshot = snapshot + self.binding_id = binding_id + self.workspace_id = workspace_id + self.backend_binding_ref = backend_binding_ref + self.resolved_scopes: list[WorkflowAgentSessionScope] = [] + self.existing_scope_lookups: list[dict[str, object]] = [] # ENG-638: set to simulate resume after a submitted/timed-out form. self.loaded_session: StoredWorkflowAgentSession | None = None self.saved: list[ @@ -141,40 +164,43 @@ class FakeSessionStore: WorkflowAgentSessionScope, str, CompositorSessionSnapshot | None, - list[RuntimeLayerSpec], str | None, str | None, ] ] = [] - self.cleaned: list[tuple[WorkflowAgentSessionScope, str | None]] = [] - def load_active_snapshot(self, scope: WorkflowAgentSessionScope) -> CompositorSessionSnapshot | None: - return self.loaded_snapshot + def load_existing_node_execution_scope(self, **kwargs: object) -> WorkflowAgentSessionScope | None: + self.existing_scope_lookups.append(kwargs) + return self.loaded_session.scope if self.loaded_session is not None else None - def load_active_session(self, scope: WorkflowAgentSessionScope) -> StoredWorkflowAgentSession | None: - return self.loaded_session + def load_or_create_node_execution_session( + self, + scope: WorkflowAgentSessionScope, + *, + home_snapshot_id: str, + ) -> StoredWorkflowAgentSession: + assert home_snapshot_id == "home-1" + self.resolved_scopes.append(scope) + if self.loaded_session is not None: + return self.loaded_session + return StoredWorkflowAgentSession( + scope=scope, + binding_id=self.binding_id, + workspace_id=self.workspace_id, + backend_binding_ref=self.backend_binding_ref, + session_snapshot=self.loaded_snapshot, + ) def save_active_snapshot( self, *, scope: WorkflowAgentSessionScope, - backend_run_id: str, + binding_id: str, snapshot: CompositorSessionSnapshot | None, - runtime_layer_specs: list[RuntimeLayerSpec], pending_form_id: str | None = None, pending_tool_call_id: str | None = None, ) -> None: - self.saved.append( - (scope, backend_run_id, snapshot, list(runtime_layer_specs), pending_form_id, pending_tool_call_id) - ) - - def mark_cleaned( - self, - *, - scope: WorkflowAgentSessionScope, - backend_run_id: str | None = None, - ) -> None: - self.cleaned.append((scope, backend_run_id)) + self.saved.append((scope, binding_id, snapshot, pending_form_id, pending_tool_call_id)) class FileOutputBackendClient(FakeAgentBackendRunClient): @@ -286,6 +312,8 @@ def _node( session_store: FakeSessionStore | None = None, declared_outputs: list[dict[str, object]] | None = None, agent_backend_client: FakeAgentBackendRunClient | None = None, + binding_resolver: FakeBindingResolver | None = None, + runtime_request_builder: WorkflowAgentRuntimeRequestBuilder | None = None, ) -> DifyAgentNode: graph_init_params = GraphInitParams( workflow_id="workflow-1", @@ -309,7 +337,7 @@ def _node( return True client = agent_backend_client or FakeAgentBackendRunClient(scenario=scenario) - binding_resolver = FakeBindingResolver() + binding_resolver = binding_resolver or FakeBindingResolver() if declared_outputs is not None: binding_resolver.binding.node_job_config = WorkflowNodeJobConfig.model_validate( { @@ -323,15 +351,22 @@ def _node( node_id="agent-node", data=DifyAgentNodeData.model_validate({"type": BuiltinNodeTypes.AGENT, "version": "2"}), graph_init_params=graph_init_params, - graph_runtime_state=cast(GraphRuntimeState, SimpleNamespace(variable_pool=FakeVariablePool())), + graph_runtime_state=cast( + GraphRuntimeState, + SimpleNamespace( + variable_pool=FakeVariablePool(), + graph_execution=SimpleNamespace(node_executions={}), + ), + ), binding_resolver=binding_resolver, - runtime_request_builder=WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()), + runtime_request_builder=runtime_request_builder + or WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()), agent_backend_client=client, event_adapter=AgentBackendRunEventAdapter(), output_adapter=WorkflowAgentOutputAdapter(), type_checker=PerOutputTypeChecker(file_validator=_AlwaysAllowFileValidator()), failure_orchestrator=OutputFailureOrchestrator(), - session_store=cast(WorkflowAgentRuntimeSessionStore | None, session_store), + session_store=cast(WorkflowAgentWorkspaceStore, session_store or FakeSessionStore()), ) @@ -367,6 +402,85 @@ def test_agent_node_run_maps_successful_agent_backend_run_to_node_result(): assert layers["llm"]["config"]["credentials"] == "[REDACTED]" +def test_agent_node_uses_resolved_backend_binding_before_backend_invocation() -> None: + client = FakeAgentBackendRunClient() + store = FakeSessionStore(binding_id="binding-2", backend_binding_ref="backend-binding-2") + + events = list(_node(agent_backend_client=client, session_store=store)._run()) + + assert len(events) == 1 + assert client.request is not None + layers = {layer["name"]: layer for layer in client.request.model_dump(mode="json")["composition"]["layers"]} + assert layers["runtime"]["config"]["backend_binding_ref"] == "backend-binding-2" + assert store.saved[0][1] == "binding-2" + assert len(store.resolved_scopes) == 1 + + +def test_agent_node_resume_resolves_the_generation_from_the_persisted_execution() -> None: + binding_resolver = FakeBindingResolver() + store = FakeSessionStore() + store.loaded_session = StoredWorkflowAgentSession( + scope=WorkflowAgentSessionScope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="workflow-run-1", + node_id="agent-node", + node_execution_id="exec-1", + workflow_agent_binding_id="binding-1", + agent_id="agent-1", + agent_config_snapshot_id="snapshot-pinned", + ), + binding_id="workspace-binding-1", + workspace_id="workspace-1", + backend_binding_ref="backend-binding-1", + session_snapshot=None, + ) + node = _node(binding_resolver=binding_resolver, session_store=store) + + events = list(node._run()) + + assert len(events) == 1 + assert store.existing_scope_lookups[0]["node_execution_id"] == node.execution_id + assert binding_resolver.calls[0]["binding_id"] == "binding-1" + assert binding_resolver.calls[0]["snapshot_id"] == "snapshot-pinned" + + +def test_agent_node_maps_persisted_participant_lookup_error_to_node_failure() -> None: + class _UnavailableParticipantStore(FakeSessionStore): + def load_existing_node_execution_scope(self, **kwargs: object) -> WorkflowAgentSessionScope | None: + del kwargs + raise AgentWorkspaceNotFoundError("Workflow node participant Binding is unavailable") + + events = list(_node(session_store=_UnavailableParticipantStore())._run()) + + assert len(events) == 1 + result = cast(StreamCompletedEvent, events[0]).node_run_result + assert result.status == WorkflowNodeExecutionStatus.FAILED + assert result.error == "Workflow node participant Binding is unavailable" + assert result.error_type == "agent_workflow_node_runtime_error" + + +def test_agent_node_passes_execution_id_to_session_store_and_runtime_request_builder() -> None: + store = FakeSessionStore() + request_builder = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()) + node = _node(session_store=store, runtime_request_builder=request_builder) + execution_id = node.ensure_execution_id() + + with patch.object(request_builder, "build", wraps=request_builder.build) as build: + list(node._run()) + + assert str(UUID(execution_id)) == execution_id + assert execution_id != node.id + scope = store.resolved_scopes[0] + build.assert_called_once() + context = cast(WorkflowAgentRuntimeBuildContext, build.call_args.args[0]) + assert scope.node_id == node.id + assert scope.node_execution_id == execution_id + assert context.node_id == node.id + assert context.node_execution_id == execution_id + + def test_agent_node_run_ignores_agent_message_delta_until_terminal_result(): events = list(_node(agent_backend_client=AgentMessageDeltaBackendClient())._run()) @@ -491,23 +605,6 @@ def test_agent_node_run_maps_failed_agent_backend_run_to_node_result(): assert result.error_type == "unit_test" -def test_agent_node_failed_run_marks_session_cleaned_to_prevent_stale_reuse(): - """A failed agent run must retire the local ACTIVE session row so a workflow - loop back into the same Agent node does not resume from a stale snapshot.""" - existing_snapshot = CompositorSessionSnapshot(layers=[]) - store = FakeSessionStore(snapshot=existing_snapshot) - - events = list(_node(scenario=FakeAgentBackendScenario.FAILED, session_store=store)._run()) - - assert len(events) == 1 - assert store.cleaned, "failed agent run should mark the session cleaned" - cleaned_scope, cleaned_backend_run_id = store.cleaned[0] - assert cleaned_scope.workflow_run_id == "workflow-run-1" - assert cleaned_backend_run_id == "fake-run-1" - # A failed run does not produce a fresh snapshot to persist. - assert store.saved == [] - - def test_agent_node_saves_success_snapshot_and_reuses_existing_snapshot(): existing_snapshot = CompositorSessionSnapshot(layers=[]) store = FakeSessionStore(snapshot=existing_snapshot) @@ -518,21 +615,15 @@ def test_agent_node_saves_success_snapshot_and_reuses_existing_snapshot(): assert len(events) == 1 assert store.saved - scope, backend_run_id, saved_snapshot, saved_specs, pending_form_id, pending_tool_call_id = store.saved[0] + scope, binding_id, saved_snapshot, pending_form_id, pending_tool_call_id = store.saved[0] assert scope.workflow_run_id == "workflow-run-1" - assert backend_run_id == "fake-run-1" + assert binding_id == "binding-1" assert saved_snapshot is not None # A successful terminal carries no ask_human pause correlation. assert pending_form_id is None assert pending_tool_call_id is None assert client.request is not None assert client.request.session_snapshot is existing_snapshot - # Persist enough composition shape to replay a cleanup run; plugin layers - # (which would carry credentials) are intentionally absent. - saved_layer_names = [spec.name for spec in saved_specs] - assert saved_layer_names, "cleanup specs must persist at least the non-plugin layers" - plugin_types = {"dify.plugin.llm", "dify.plugin.tools"} - assert not {spec.type for spec in saved_specs} & plugin_types def test_agent_node_run_when_session_store_save_raises_records_persist_error_in_metadata(): @@ -553,83 +644,7 @@ def test_agent_node_run_when_session_store_save_raises_records_persist_error_in_ assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED agent_backend = result.metadata[WorkflowNodeExecutionMetadataKey.AGENT_LOG]["agent_backend"] assert agent_backend["session_snapshot_persisted"] is False - assert agent_backend["session_snapshot_persist_error"] == "workflow_agent_runtime_session_store_error" - - -def test_agent_node_failed_run_when_mark_cleaned_raises_records_cleanup_error_in_metadata(): - """Same defensive pattern: a DB-side mark_cleaned failure must surface as - a ``session_snapshot_cleanup_error`` in metadata, not as a node crash.""" - - class _ExplodingMarkCleanedStore(FakeSessionStore): - def mark_cleaned(self, **kwargs): # type: ignore[override] - del kwargs - raise RuntimeError("simulated DB failure") - - store = _ExplodingMarkCleanedStore() - events = list(_node(scenario=FakeAgentBackendScenario.FAILED, session_store=store)._run()) - - assert len(events) == 1 - result = cast(StreamCompletedEvent, events[0]).node_run_result - assert result.status == WorkflowNodeExecutionStatus.FAILED - agent_backend = result.metadata[WorkflowNodeExecutionMetadataKey.AGENT_LOG]["agent_backend"] - assert agent_backend["session_snapshot_cleaned_on_failure"] is False - assert agent_backend["session_snapshot_cleanup_error"] == "workflow_agent_runtime_session_store_error" - - -def test_agent_node_success_run_without_session_store_skips_persistence(): - """When ``session_store`` is None the node still completes successfully — - the lifecycle branch is a no-op and the run result is unaffected.""" - events = list(_node(session_store=None)._run()) - - assert len(events) == 1 - result = cast(StreamCompletedEvent, events[0]).node_run_result - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - agent_backend = result.metadata[WorkflowNodeExecutionMetadataKey.AGENT_LOG]["agent_backend"] - # No persistence metadata is attached when the store is missing. - assert "session_snapshot_persisted" not in agent_backend - - -def test_agent_node_failed_run_without_session_store_skips_mark_cleaned(): - """``session_store=None`` + failed terminal must remain a no-op for - the cleanup branch — the node failure path still surfaces correctly.""" - events = list(_node(scenario=FakeAgentBackendScenario.FAILED, session_store=None)._run()) - - assert len(events) == 1 - result = cast(StreamCompletedEvent, events[0]).node_run_result - assert result.status == WorkflowNodeExecutionStatus.FAILED - agent_backend = result.metadata[WorkflowNodeExecutionMetadataKey.AGENT_LOG]["agent_backend"] - assert "session_snapshot_cleaned_on_failure" not in agent_backend - - -def test_agent_node_failed_run_enqueues_backend_cleanup_before_local_retirement(monkeypatch): - store = FakeSessionStore() - store.loaded_session = StoredWorkflowAgentSession( - scope=_pending_session(CompositorSessionSnapshot(layers=[])).scope, - session_snapshot=CompositorSessionSnapshot(layers=[]), - backend_run_id="stored-run-1", - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - ) - queued_payloads: list[dict[str, object]] = [] - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.agent_node.cleanup_workflow_agent_runtime_session.delay", - lambda payload: queued_payloads.append(payload), - ) - - events = list(_node(scenario=FakeAgentBackendScenario.FAILED, session_store=store)._run()) - - assert len(events) == 1 - result = cast(StreamCompletedEvent, events[0]).node_run_result - assert result.status == WorkflowNodeExecutionStatus.FAILED - assert store.cleaned[0][1] == "fake-run-1" - assert store.cleaned[0][0].workflow_run_id == "workflow-run-1" - assert store.cleaned[0][0].node_id == "agent-node" - assert len(queued_payloads) == 1 - assert ( - queued_payloads[0]["idempotency_key"] - == "tenant-1:workflow-run-1:agent-node:binding-1:workflow-agent-failure-cleanup:stored-run-1:fake-run-1" - ) - assert queued_payloads[0]["metadata"]["previous_agent_backend_run_id"] == "stored-run-1" - assert queued_payloads[0]["metadata"]["failed_agent_backend_run_id"] == "fake-run-1" + assert agent_backend["session_snapshot_persist_error"] == "workflow_agent_workspace_store_error" def test_agent_node_paused_run_requests_workflow_pause_and_persists_snapshot(): @@ -646,17 +661,21 @@ def test_agent_node_paused_run_requests_workflow_pause_and_persists_snapshot(): events = list(node._run()) assert len(events) == 1 - assert isinstance(events[0], PauseRequestedEvent) + assert isinstance(events[0], NodeRunPauseRequestedEvent) assert isinstance(events[0].reason, HitlRequired) assert events[0].reason.session_id == "form-1" assert events[0].reason.node_id == "agent-node" + assert events[0].node_run_result.process_data == { + "agent_id": "agent-1", + "agent_config_snapshot_id": "snapshot-1", + "workflow_agent_binding_id": "binding-1", + } fake_repo.create_form.assert_called_once() assert store.saved - assert store.saved[0][1] == "fake-run-1" - assert store.saved[0][3], "paused agent run should still persist replayable layer specs" + assert store.saved[0][1] == "binding-1" # ENG-637: the awaiting form + deferred tool_call correlation is persisted. - assert store.saved[0][4] == "form-1" - assert store.saved[0][5] == "fake-ask-human-1" + assert store.saved[0][3] == "form-1" + assert store.saved[0][4] == "fake-ask-human-1" def _pending_session(snapshot: CompositorSessionSnapshot) -> StoredWorkflowAgentSession: @@ -668,12 +687,14 @@ def _pending_session(snapshot: CompositorSessionSnapshot) -> StoredWorkflowAgent workflow_run_id="workflow-run-1", node_id="agent-node", node_execution_id="exec-1", - binding_id="binding-1", + workflow_agent_binding_id="binding-1", agent_id="agent-1", agent_config_snapshot_id="snapshot-1", ), + binding_id="binding-1", + workspace_id="workspace-1", + backend_binding_ref="backend-binding-1", session_snapshot=snapshot, - backend_run_id="run-0", pending_form_id="form-1", pending_tool_call_id="call-1", ) @@ -727,19 +748,75 @@ def test_agent_node_repauses_when_resumed_form_still_waiting(monkeypatch): events = list(node._run()) assert len(events) == 1 - assert isinstance(events[0], PauseRequestedEvent) + assert isinstance(events[0], NodeRunPauseRequestedEvent) assert isinstance(events[0].reason, HitlRequired) + assert events[0].node_run_result.process_data["workflow_agent_binding_id"] == "binding-1" assert client.request is None # no second Agent run was created +def test_agent_node_expired_ask_human_failure_keeps_binding_identity(monkeypatch): + snapshot = CompositorSessionSnapshot(layers=[]) + store = FakeSessionStore(snapshot=snapshot) + store.loaded_session = _pending_session(snapshot) + + def _raise_expired_form(**_kwargs): + raise AssertionError("cannot resume globally expired ask_human form, form_id=form-1") + + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.agent_node.resolve_ask_human_form", + _raise_expired_form, + ) + + events = list(_node(session_store=store)._run()) + + assert len(events) == 1 + result = cast(StreamCompletedEvent, events[0]).node_run_result + assert result.status == WorkflowNodeExecutionStatus.FAILED + assert result.error == "cannot resume globally expired ask_human form, form_id=form-1" + assert result.error_type == "agent_workflow_node_runtime_error" + assert result.process_data["workflow_agent_binding_id"] == "binding-1" + assert "agent_workspace_binding_id" not in result.process_data + + +def test_agent_node_unexpected_post_resolution_failure_keeps_binding_identity(): + class _FailingSessionStore(FakeSessionStore): + def load_or_create_node_execution_session( + self, + scope: WorkflowAgentSessionScope, + *, + home_snapshot_id: str, + ) -> StoredWorkflowAgentSession: + del scope, home_snapshot_id + raise RuntimeError("session store failed") + + events = list(_node(session_store=_FailingSessionStore())._run()) + + assert len(events) == 1 + result = cast(StreamCompletedEvent, events[0]).node_run_result + assert result.status == WorkflowNodeExecutionStatus.FAILED + assert result.error == "session store failed" + assert result.error_type == "agent_workflow_node_runtime_error" + assert result.process_data == { + "agent_id": "agent-1", + "agent_config_snapshot_id": "snapshot-1", + "workflow_agent_binding_id": "binding-1", + } + + def test_agent_node_cancels_backend_run_when_stream_fails(): client = FailingStreamBackendClient() node = _node(agent_backend_client=client) - terminal, failure = node._consume_event_stream("run-1", {"agent_backend": {}}) + terminal, failure = node._consume_event_stream( + "run-1", + inputs={}, + process_data={"workflow_agent_binding_id": "binding-1"}, + metadata={"agent_backend": {}}, + ) assert terminal is None assert failure is not None + assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"} assert len(client.cancel_requests) == 1 assert client.cancel_requests[0] is not None assert client.cancel_requests[0].reason == "event_stream_failed" @@ -749,7 +826,12 @@ def test_agent_node_cancels_backend_run_when_stream_ends_without_terminal_event( client = EmptyStreamBackendClient() node = _node(agent_backend_client=client) - terminal, failure = node._consume_event_stream("run-1", {"agent_backend": {}}) + terminal, failure = node._consume_event_stream( + "run-1", + inputs={}, + process_data={"workflow_agent_binding_id": "binding-1"}, + metadata={"agent_backend": {}}, + ) assert terminal is None assert failure is None @@ -761,11 +843,17 @@ def test_agent_node_cancels_backend_run_when_stream_raises_unexpected_error(): client = GenericFailingStreamBackendClient() node = _node(agent_backend_client=client) - terminal, failure = node._consume_event_stream("run-1", {"agent_backend": {}}) + terminal, failure = node._consume_event_stream( + "run-1", + inputs={}, + process_data={"workflow_agent_binding_id": "binding-1"}, + metadata={"agent_backend": {}}, + ) assert terminal is None assert failure is not None assert failure.node_run_result.error == "unexpected stream failure" + assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"} assert client.cancel_requests[0] is not None assert client.cancel_requests[0].reason == "event_stream_failed" @@ -775,7 +863,12 @@ def test_agent_node_uses_graph_abort_reason_when_cancel_request_fails(caplog): node = _node(agent_backend_client=client) node.graph_runtime_state.graph_execution = SimpleNamespace(aborted=True) - terminal, failure = node._consume_event_stream("run-1", {"agent_backend": {}}) + terminal, failure = node._consume_event_stream( + "run-1", + inputs={}, + process_data={"workflow_agent_binding_id": "binding-1"}, + metadata={"agent_backend": {}}, + ) assert terminal is None assert failure is not None @@ -794,13 +887,19 @@ def test_agent_node_cancels_backend_run_for_unexpected_internal_event(): return_value=[SimpleNamespace(type=AgentBackendInternalEventType.RUN_FAILED)] ) - terminal, failure = node._consume_event_stream("run-1", {"agent_backend": {}}) + terminal, failure = node._consume_event_stream( + "run-1", + inputs={}, + process_data={"workflow_agent_binding_id": "binding-1"}, + metadata={"agent_backend": {}}, + ) assert terminal is None assert failure is not None assert failure.node_run_result.error == ( "Unexpected internal event type " ) + assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"} node._agent_backend_client.cancel_run.assert_called_once() diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_binding_resolver.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_binding_resolver.py index 86f3046d47f..d52abb46128 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_binding_resolver.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_binding_resolver.py @@ -23,6 +23,7 @@ class FakeSession: def __init__(self, scalar_results): self._scalar_results = list(scalar_results) self.expunge_calls = [] + self.scalar_statements = [] def __enter__(self): return self @@ -30,7 +31,8 @@ class FakeSession: def __exit__(self, exc_type, exc, tb): return False - def scalar(self, _stmt): + def scalar(self, stmt): + self.scalar_statements.append(stmt) if not self._scalar_results: return None return self._scalar_results.pop(0) @@ -62,6 +64,7 @@ def _snapshot() -> AgentConfigSnapshot: tenant_id="tenant-1", agent_id="agent-1", version=1, + home_snapshot_id="home-1", config_snapshot=AgentSoulConfig( model=AgentSoulModelConfig( plugin_id="langgenius/openai", @@ -115,6 +118,73 @@ def test_binding_resolver_uses_active_snapshot_for_roster_agent(monkeypatch: pyt assert bundle.snapshot.id == "active-snapshot" +@pytest.mark.parametrize( + "binding_type", + [WorkflowAgentBindingType.ROSTER_AGENT, WorkflowAgentBindingType.INLINE_AGENT], +) +def test_binding_resolver_uses_pinned_snapshot_for_existing_node_execution( + monkeypatch: pytest.MonkeyPatch, + binding_type: WorkflowAgentBindingType, +): + binding = _binding() + binding.binding_type = binding_type + agent = _agent() + agent.active_config_snapshot_id = "active-snapshot" + snapshot = _snapshot() + snapshot.id = "pinned-snapshot" + fake_session = FakeSession([binding, agent, snapshot]) + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.binding_resolver.session_factory.create_session", + lambda: fake_session, + ) + + bundle = WorkflowAgentBindingResolver().resolve( + **_resolve(), + binding_id="binding-1", + snapshot_id="pinned-snapshot", + ) + + assert bundle.snapshot.id == "pinned-snapshot" + assert "binding-1" in fake_session.scalar_statements[0].compile().params.values() + assert "pinned-snapshot" in fake_session.scalar_statements[-1].compile().params.values() + + +def test_binding_resolver_does_not_fallback_from_an_explicit_empty_snapshot(monkeypatch: pytest.MonkeyPatch): + binding = _binding() + binding.binding_type = WorkflowAgentBindingType.ROSTER_AGENT + agent = _agent() + agent.active_config_snapshot_id = "active-snapshot" + fake_session = FakeSession([binding, agent, None]) + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.binding_resolver.session_factory.create_session", + lambda: fake_session, + ) + + with pytest.raises(WorkflowAgentBindingError) as exc_info: + WorkflowAgentBindingResolver().resolve(**_resolve(), binding_id="binding-1", snapshot_id="") + + assert exc_info.value.error_code == "agent_config_snapshot_not_found" + assert "" in fake_session.scalar_statements[-1].compile().params.values() + + +@pytest.mark.parametrize( + ("binding_id", "snapshot_id"), + [("binding-1", None), (None, "snapshot-1")], +) +def test_binding_resolver_rejects_half_pinned_generation( + binding_id: str | None, + snapshot_id: str | None, +) -> None: + with pytest.raises(WorkflowAgentBindingError) as exc_info: + WorkflowAgentBindingResolver().resolve( + **_resolve(), + binding_id=binding_id, + snapshot_id=snapshot_id, + ) + + assert exc_info.value.error_code == "agent_binding_generation_invalid" + + def test_binding_resolver_rejects_unpublished_roster_agent(monkeypatch: pytest.MonkeyPatch): binding = _binding() binding.binding_type = WorkflowAgentBindingType.ROSTER_AGENT diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py index c66bcf5b6e6..f52bc63bdc1 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py @@ -203,6 +203,8 @@ def _context() -> WorkflowAgentRuntimeBuildContext: binding=binding, agent=agent, snapshot=snapshot, + binding_id="binding-1", + backend_binding_ref="binding-ref-1", ) @@ -467,7 +469,7 @@ def test_build_maps_agent_soul_shell_settings_to_shell_layer(monkeypatch: pytest shell_config = {layer["name"]: layer for layer in dumped["composition"]["layers"]}[DIFY_SHELL_LAYER_ID]["config"] assert shell_config["cli_tools"][0]["install_commands"] == ["apt-get install -y ripgrep"] assert shell_config["env"][0] == {"name": "PROJECT_NAME", "value": "demo"} - assert shell_config["sandbox"] == {"provider": "independent", "config": {"cpu": 2}} + assert "sandbox" not in shell_config assert result.metadata["agent_tools"] == { "dify_tool_count": 0, "dify_tool_names": [], @@ -520,7 +522,7 @@ def test_build_shell_layer_config_accepts_legacy_fallback_keys(): {"name": "API_KEY", "ref": "credential-2"}, {"name": "LEGACY_SECRET_REF", "ref": "credential-3"}, ] - assert config["sandbox"] is None + assert "sandbox" not in config def test_build_shell_layer_config_maps_typed_command_field(): @@ -1459,7 +1461,10 @@ def test_workflow_run_request_has_config_layer_with_empty_agent_soul(monkeypatch "mentioned_skill_names": [], "mentioned_file_names": [], } - assert layers[DIFY_SHELL_LAYER_ID]["deps"] == {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID} + assert layers[DIFY_SHELL_LAYER_ID]["deps"] == { + "execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID, + "runtime": "runtime", + } assert layers[DIFY_SHELL_LAYER_ID]["config"]["agent_stub_drive_ref"] is None @@ -1474,7 +1479,7 @@ def test_workflow_run_request_contains_config_layer(): layer_names = [layer["name"] for layer in dumped["composition"]["layers"]] assert DIFY_CONFIG_LAYER_ID in layer_names # shell enters first; config uses that shell to materialize mentioned targets. - assert layer_names.index(DIFY_SHELL_LAYER_ID) == layer_names.index("execution_context") + 1 + assert layer_names.index(DIFY_SHELL_LAYER_ID) == layer_names.index("execution_context") + 2 assert layer_names.index(DIFY_CONFIG_LAYER_ID) == layer_names.index(DIFY_SHELL_LAYER_ID) + 1 config = next(layer for layer in dumped["composition"]["layers"] if layer["name"] == DIFY_CONFIG_LAYER_ID) assert config["type"] == "dify.config" @@ -1498,11 +1503,6 @@ def test_workflow_run_request_contains_config_layer(): } warnings = result.metadata["runtime_support"]["unsupported_runtime_warnings"] assert warnings == [] - # the config layer is non-sensitive and must survive into persistable specs - from dify_agent.protocol import extract_runtime_layer_specs - - specs = extract_runtime_layer_specs(result.request.composition) - assert any(spec.name == DIFY_CONFIG_LAYER_ID and spec.type == "dify.config" for spec in specs) def test_workflow_runtime_expands_config_mentions_in_agent_soul_prompt(): diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py deleted file mode 100644 index eda2962198e..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py +++ /dev/null @@ -1,256 +0,0 @@ -from typing import cast - -import pytest -from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from dify_agent.protocol import RuntimeLayerSpec - -from core.workflow.nodes.agent_v2.session_cleanup_layer import ( - WorkflowAgentSessionCleanupLayer, - build_workflow_agent_session_cleanup_layer, -) -from core.workflow.nodes.agent_v2.session_store import ( - StoredWorkflowAgentSession, - WorkflowAgentRuntimeSessionStore, - WorkflowAgentSessionScope, -) -from core.workflow.system_variables import build_system_variables -from graphon.entities.pause_reason import SchedulingPause -from graphon.graph_engine.command_channels import CommandChannel -from graphon.graph_events import ( - GraphRunAbortedEvent, - GraphRunFailedEvent, - GraphRunPartialSucceededEvent, - GraphRunPausedEvent, - GraphRunStartedEvent, - GraphRunSucceededEvent, -) -from graphon.runtime import GraphRuntimeState, ReadOnlyGraphRuntimeStateWrapper, VariablePool - - -def _layer_snapshot(name: str) -> LayerSessionSnapshot: - return LayerSessionSnapshot( - name=name, - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ) - - -def _default_scope() -> WorkflowAgentSessionScope: - return WorkflowAgentSessionScope( - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - workflow_run_id="workflow-run-1", - node_id="agent-node", - node_execution_id="node-exec-1", - binding_id="binding-1", - agent_id="agent-1", - agent_config_snapshot_id="snapshot-1", - ) - - -def _stored_session(scope: WorkflowAgentSessionScope, *, index: int = 1) -> StoredWorkflowAgentSession: - return StoredWorkflowAgentSession( - scope=scope, - session_snapshot=CompositorSessionSnapshot( - layers=[ - _layer_snapshot("workflow_node_job_prompt"), - _layer_snapshot("execution_context"), - _layer_snapshot("history"), - _layer_snapshot("llm"), - ] - ), - backend_run_id=f"agent-run-{index}", - runtime_layer_specs=[ - RuntimeLayerSpec(name="workflow_node_job_prompt", type="plain.prompt", config={"prefix": "ok"}), - RuntimeLayerSpec(name="execution_context", type="dify.execution_context", config={"tenant_id": "t"}), - RuntimeLayerSpec(name="history", type="pydantic_ai.history"), - ], - ) - - -class FakeSessionStore: - def __init__(self, *, stored: list[StoredWorkflowAgentSession] | None = None) -> None: - self._stored = stored if stored is not None else [_stored_session(_default_scope())] - self.list_calls: list[str] = [] - self.cleaned: list[tuple[WorkflowAgentSessionScope, str | None]] = [] - - def list_active_sessions(self, *, workflow_run_id: str) -> list[StoredWorkflowAgentSession]: - self.list_calls.append(workflow_run_id) - return list(self._stored) - - def mark_cleaned(self, *, scope: WorkflowAgentSessionScope, backend_run_id: str | None = None) -> None: - self.cleaned.append((scope, backend_run_id)) - - -def _build_layer(*, session_store: FakeSessionStore) -> WorkflowAgentSessionCleanupLayer: - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="workflow-run-1"), - user_inputs={}, - conversation_variables=[], - ) - runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=0.0) - layer = WorkflowAgentSessionCleanupLayer( - session_store=cast(WorkflowAgentRuntimeSessionStore, session_store), - ) - layer.initialize(ReadOnlyGraphRuntimeStateWrapper(runtime_state), cast(CommandChannel, object())) - return layer - - -@pytest.mark.parametrize( - "terminal_event", - [ - GraphRunSucceededEvent(outputs={}), - GraphRunPartialSucceededEvent(exceptions_count=1, outputs={}), - GraphRunFailedEvent(error="boom"), - GraphRunAbortedEvent(reason="user cancelled", outputs={}), - ], - ids=["succeeded", "partial_succeeded", "failed", "aborted"], -) -def test_cleanup_layer_enqueues_cleanup_and_marks_cleaned_on_terminal_events(monkeypatch, terminal_event): - session_store = FakeSessionStore() - queued_payloads: list[dict[str, object]] = [] - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", - lambda payload: queued_payloads.append(payload), - ) - layer = _build_layer(session_store=session_store) - - layer.on_event(terminal_event) - - assert session_store.list_calls == ["workflow-run-1"] - assert len(queued_payloads) == 1 - assert queued_payloads[0]["metadata"]["workflow_run_id"] == "workflow-run-1" - assert queued_payloads[0]["metadata"]["previous_agent_backend_run_id"] == "agent-run-1" - assert session_store.cleaned == [(_default_scope(), "agent-run-1")] - - -@pytest.mark.parametrize( - "non_terminal_event", - [ - GraphRunStartedEvent(), - GraphRunPausedEvent(reasons=[SchedulingPause(message="awaiting human input")], outputs={}), - ], - ids=["started", "paused"], -) -def test_cleanup_layer_ignores_non_terminal_events(monkeypatch, non_terminal_event): - session_store = FakeSessionStore() - queued_payloads: list[dict[str, object]] = [] - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", - lambda payload: queued_payloads.append(payload), - ) - layer = _build_layer(session_store=session_store) - - layer.on_event(non_terminal_event) - - assert session_store.list_calls == [] - assert queued_payloads == [] - assert session_store.cleaned == [] - - -def test_cleanup_layer_marks_cleaned_even_when_specs_are_missing(monkeypatch, caplog: pytest.LogCaptureFixture): - scope = _default_scope() - session_store = FakeSessionStore( - stored=[ - StoredWorkflowAgentSession( - scope=scope, - session_snapshot=CompositorSessionSnapshot(layers=[_layer_snapshot("history")]), - backend_run_id="legacy-run", - runtime_layer_specs=[], - ) - ] - ) - queued_payloads: list[dict[str, object]] = [] - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", - lambda payload: queued_payloads.append(payload), - ) - layer = _build_layer(session_store=session_store) - - layer.on_event(GraphRunSucceededEvent(outputs={})) - - assert queued_payloads == [] - assert session_store.cleaned == [(scope, "legacy-run")] - assert any("no runtime_layer_specs persisted" in record.message for record in caplog.records) - - -def test_cleanup_layer_marks_cleaned_even_when_enqueue_fails(monkeypatch): - session_store = FakeSessionStore() - - def _explode(_payload: dict[str, object]) -> None: - raise RuntimeError("queue down") - - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", - _explode, - ) - layer = _build_layer(session_store=session_store) - - layer.on_event(GraphRunSucceededEvent(outputs={})) - - assert session_store.cleaned == [(_default_scope(), "agent-run-1")] - - -def test_cleanup_layer_does_not_raise_when_mark_cleaned_fails(monkeypatch): - session_store = FakeSessionStore() - - def _explode(*, scope: WorkflowAgentSessionScope, backend_run_id: str | None = None) -> None: - del scope, backend_run_id - raise RuntimeError("cleanup bookkeeping failed") - - monkeypatch.setattr(session_store, "mark_cleaned", _explode) - layer = _build_layer(session_store=session_store) - - layer.on_event(GraphRunSucceededEvent(outputs={})) - - -def test_cleanup_layer_fans_out_to_every_active_session(monkeypatch): - scopes = [ - WorkflowAgentSessionScope( - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - workflow_run_id="workflow-run-1", - node_id=f"agent-node-{i}", - node_execution_id=f"node-exec-{i}", - binding_id=f"binding-{i}", - agent_id=f"agent-{i}", - agent_config_snapshot_id=f"snapshot-{i}", - ) - for i in range(3) - ] - session_store = FakeSessionStore(stored=[_stored_session(scope, index=i) for i, scope in enumerate(scopes, 1)]) - queued_payloads: list[dict[str, object]] = [] - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", - lambda payload: queued_payloads.append(payload), - ) - layer = _build_layer(session_store=session_store) - - layer.on_event(GraphRunSucceededEvent(outputs={})) - - assert len(queued_payloads) == 3 - assert [entry[0] for entry in session_store.cleaned] == scopes - - -def test_cleanup_layer_skips_when_workflow_run_id_missing(caplog: pytest.LogCaptureFixture): - session_store = FakeSessionStore() - variable_pool = VariablePool.from_bootstrap(system_variables={}, user_inputs={}, conversation_variables=[]) - runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=0.0) - layer = WorkflowAgentSessionCleanupLayer(session_store=cast(WorkflowAgentRuntimeSessionStore, session_store)) - layer.initialize(ReadOnlyGraphRuntimeStateWrapper(runtime_state), cast(CommandChannel, object())) - - layer.on_event(GraphRunSucceededEvent(outputs={})) - - assert session_store.list_calls == [] - assert session_store.cleaned == [] - assert any("workflow_run_id is missing" in record.message for record in caplog.records) - - -def test_build_workflow_agent_session_cleanup_layer_returns_layer() -> None: - layer = build_workflow_agent_session_cleanup_layer() - - assert isinstance(layer, WorkflowAgentSessionCleanupLayer) diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_store.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_store.py index a57b2adc168..3d3edd58137 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_store.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_store.py @@ -1,288 +1,428 @@ -"""Unit tests for :mod:`core.workflow.nodes.agent_v2.session_store`. - -Uses the in-memory SQLite engine configured by the project conftest plus a -per-test ``CREATE TABLE`` so the real ORM round-trip exercises every store -method. Keeps the suite self-contained — no Postgres / Docker required — while -still hitting the actual ``session_factory`` code path that production uses. -""" - -from __future__ import annotations - -from collections.abc import Generator +import json +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock import pytest from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from dify_agent.protocol import RuntimeLayerSpec -from sqlalchemy import delete +from sqlalchemy.orm import Session -from core.db.session_factory import session_factory -from core.workflow.nodes.agent_v2.session_store import ( - StoredWorkflowAgentSession, - WorkflowAgentRuntimeSessionStore, - WorkflowAgentSessionScope, +from core.workflow.nodes.agent_v2.session_store import WorkflowAgentSessionScope, WorkflowAgentWorkspaceStore +from models.agent import ( + AgentConfigVersionKind, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, ) -from models.agent import WorkflowAgentRuntimeSession, WorkflowAgentRuntimeSessionStatus +from models.workflow import WorkflowNodeExecutionModel +from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService +from services.agent_app_sandbox_service import WorkflowAgentSandboxService -def _scope(workflow_run_id: str | None = "wfr-1", binding_id: str = "binding-1") -> WorkflowAgentSessionScope: +def _scope() -> WorkflowAgentSessionScope: return WorkflowAgentSessionScope( tenant_id="tenant-1", app_id="app-1", workflow_id="workflow-1", - workflow_run_id=workflow_run_id, - node_id="agent-node", - node_execution_id="node-exec-1", - binding_id=binding_id, + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + workflow_agent_binding_id="workflow-binding-1", agent_id="agent-1", - agent_config_snapshot_id="snapshot-1", + agent_config_snapshot_id="config-1", ) -def _snapshot(messages: int = 1) -> CompositorSessionSnapshot: - return CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="history", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={"messages": [{"role": "user", "content": f"m{i}"} for i in range(messages)]}, - ) - ] +def _binding() -> SimpleNamespace: + return SimpleNamespace( + id="binding-1", + workspace_id="workspace-1", + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="backend-binding-1", + session_snapshot=None, + pending_form_id=None, + pending_tool_call_id=None, ) -def _specs() -> list[RuntimeLayerSpec]: - return [ - RuntimeLayerSpec(name="workflow_node_job_prompt", type="plain.prompt", config={"prefix": "ok"}), - RuntimeLayerSpec(name="history", type="pydantic_ai.history"), - ] - - -@pytest.fixture(autouse=True) -def _create_table() -> Generator[None, None, None]: - """Create the lifecycle table on the in-memory SQLite engine, drop after.""" - engine = session_factory.get_session_maker().kw["bind"] - WorkflowAgentRuntimeSession.__table__.create(bind=engine, checkfirst=True) - yield - with session_factory.create_session() as session: - session.execute(delete(WorkflowAgentRuntimeSession)) - session.commit() - WorkflowAgentRuntimeSession.__table__.drop(bind=engine, checkfirst=True) - - -def test_load_active_snapshot_returns_none_when_scope_has_no_workflow_run_id(): - """``workflow_run_id`` is the keying column; no row can match without it.""" - store = WorkflowAgentRuntimeSessionStore() - assert store.load_active_snapshot(_scope(workflow_run_id=None)) is None - - -def test_load_active_snapshot_returns_none_when_no_row_matches(): - store = WorkflowAgentRuntimeSessionStore() - assert store.load_active_snapshot(_scope()) is None - - -def test_save_active_snapshot_creates_row_and_load_round_trips(): - store = WorkflowAgentRuntimeSessionStore() - snapshot = _snapshot(messages=2) - store.save_active_snapshot(scope=_scope(), backend_run_id="run-1", snapshot=snapshot, runtime_layer_specs=_specs()) - - loaded = store.load_active_snapshot(_scope()) - assert loaded is not None - assert len(loaded.layers) == 1 - assert loaded.layers[0].name == "history" - assert loaded.layers[0].runtime_state["messages"] == snapshot.layers[0].runtime_state["messages"] - with session_factory.create_session() as session: - row = session.query(WorkflowAgentRuntimeSession).one() - assert "workflow_node_job_prompt" in row.composition_layer_specs - assert "history" in row.composition_layer_specs - - -def test_save_active_snapshot_skips_when_workflow_run_id_missing(): - """Without a workflow_run_id the row cannot be keyed; save is a no-op.""" - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(workflow_run_id=None), - backend_run_id="run-skipped", - snapshot=_snapshot(), - runtime_layer_specs=_specs(), - ) - with session_factory.create_session() as session: - assert session.query(WorkflowAgentRuntimeSession).count() == 0 - - -def test_save_active_snapshot_skips_when_snapshot_missing(): - """A run that produced no snapshot (e.g. failed agent run) does not write.""" - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-empty", - snapshot=None, - runtime_layer_specs=_specs(), - ) - with session_factory.create_session() as session: - assert session.query(WorkflowAgentRuntimeSession).count() == 0 - - -def test_save_active_snapshot_updates_existing_row_on_re_entry(): - """A second save under the same scope must update in place, not insert.""" - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(messages=1), - runtime_layer_specs=_specs(), - ) - # Second call with new snapshot + backend_run_id. - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-2", - snapshot=_snapshot(messages=2), - runtime_layer_specs=_specs(), +def _workspace_row( + *, + workspace_id: str = "workspace-1", + tenant_id: str = "tenant-1", + app_id: str = "app-1", + owner_scope_key: str = "node-1:workflow-binding-1", + status: AgentWorkingResourceStatus = AgentWorkingResourceStatus.ACTIVE, +) -> AgentWorkspace: + return AgentWorkspace( + id=workspace_id, + tenant_id=tenant_id, + app_id=app_id, + owner_type=AgentWorkspaceOwnerType.WORKFLOW_RUN, + owner_id="run-1", + owner_scope_key=owner_scope_key, + backend_workspace_ref=f"{workspace_id}-ref", + status=status, + active_guard=1 if status is AgentWorkingResourceStatus.ACTIVE else None, ) - with session_factory.create_session() as session: - rows = session.query(WorkflowAgentRuntimeSession).all() - assert len(rows) == 1 - assert rows[0].backend_run_id == "run-2" - assert rows[0].status == WorkflowAgentRuntimeSessionStatus.ACTIVE - assert rows[0].cleaned_at is None - -def test_save_active_snapshot_resurrects_cleaned_row(): - """If a prior cleanup retired the row, a re-entry flips it back to ACTIVE.""" - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(), - runtime_layer_specs=_specs(), - ) - store.mark_cleaned(scope=_scope(), backend_run_id="cleanup-1") - # Save again — the existing row was CLEANED; should be revived. - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-2", - snapshot=_snapshot(messages=3), - runtime_layer_specs=_specs(), +def _binding_row( + *, + binding_id: str = "binding-1", + workspace_id: str = "workspace-1", + status: AgentWorkingResourceStatus = AgentWorkingResourceStatus.ACTIVE, +) -> AgentWorkspaceBinding: + return AgentWorkspaceBinding( + id=binding_id, + tenant_id="tenant-1", + app_id="app-1", + workspace_id=workspace_id, + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref=f"{binding_id}-ref", + status=status, ) - with session_factory.create_session() as session: - rows = session.query(WorkflowAgentRuntimeSession).all() - assert len(rows) == 1 - assert rows[0].status == WorkflowAgentRuntimeSessionStatus.ACTIVE - assert rows[0].cleaned_at is None - assert rows[0].backend_run_id == "run-2" + +def test_scope_uses_node_and_workflow_binding_as_workspace_subscope() -> None: + owner = _scope().workspace_owner + assert owner.owner_type is AgentWorkspaceOwnerType.WORKFLOW_RUN + assert owner.owner_id == "run-1" + assert owner.owner_scope_key == "node-1:workflow-binding-1" -def test_list_active_sessions_returns_specs_and_snapshot(): - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(binding_id="binding-A"), - backend_run_id="run-A", - snapshot=_snapshot(), - runtime_layer_specs=_specs(), +def test_load_existing_scope_reads_the_generation_from_the_persisted_binding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + execution = SimpleNamespace( + agent_workspace_binding_id="binding-1", + process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"}, ) - store.save_active_snapshot( - scope=_scope(binding_id="binding-B"), - backend_run_id="run-B", - snapshot=_snapshot(messages=2), - runtime_layer_specs=_specs(), + context = MagicMock() + store = WorkflowAgentWorkspaceStore() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr(store, "_load_execution_by_identity", MagicMock(return_value=execution)) + get_active = MagicMock(return_value=_binding()) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active) + + scope = store.load_existing_node_execution_scope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", ) - listed = store.list_active_sessions(workflow_run_id="wfr-1") - assert {s.backend_run_id for s in listed} == {"run-A", "run-B"} - by_run = {s.backend_run_id: s for s in listed} - assert isinstance(by_run["run-A"], StoredWorkflowAgentSession) - # Specs round-trip through pydantic TypeAdapter — ensure deserialize works. - assert by_run["run-A"].runtime_layer_specs[0].name == "workflow_node_job_prompt" - assert by_run["run-A"].runtime_layer_specs[1].type == "pydantic_ai.history" - # node_execution_id default-replaces NULL with "" when the DB column is None. - assert by_run["run-A"].scope.node_execution_id == "node-exec-1" + assert scope is not None + assert scope.workflow_agent_binding_id == "workflow-binding-1" + assert scope.agent_id == "agent-1" + assert scope.agent_config_snapshot_id == "config-1" + assert get_active.call_args.kwargs["binding_id"] == "binding-1" -def test_list_active_sessions_skips_cleaned_rows(): - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(binding_id="binding-A"), - backend_run_id="run-A", - snapshot=_snapshot(), - runtime_layer_specs=_specs(), +def test_load_existing_scope_rejects_unavailable_persisted_binding(monkeypatch: pytest.MonkeyPatch) -> None: + execution = SimpleNamespace( + agent_workspace_binding_id="binding-missing", + process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"}, ) - store.save_active_snapshot( - scope=_scope(binding_id="binding-B"), - backend_run_id="run-B", - snapshot=_snapshot(), - runtime_layer_specs=_specs(), + context = MagicMock() + store = WorkflowAgentWorkspaceStore() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, ) - store.mark_cleaned(scope=_scope(binding_id="binding-A"), backend_run_id="cleanup-A") + monkeypatch.setattr(store, "_load_execution_by_identity", MagicMock(return_value=execution)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None)) - listed = store.list_active_sessions(workflow_run_id="wfr-1") - assert {s.backend_run_id for s in listed} == {"run-B"} + with pytest.raises(AgentWorkspaceNotFoundError, match="participant Binding is unavailable"): + store.load_existing_node_execution_scope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + ) -def test_list_active_sessions_handles_legacy_rows_without_specs(): - """Rows persisted before runtime_layer_specs landed have an empty string.""" - # Insert a legacy-shape row directly: empty specs payload simulates a row - # written before the spec persistence feature landed in A.1. - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-legacy", - snapshot=_snapshot(), - runtime_layer_specs=[], +def test_load_or_create_persists_binding_on_node_execution(monkeypatch) -> None: + execution = WorkflowNodeExecutionModel( + agent_workspace_binding_id=None, + process_data=json.dumps({"existing": "value"}), ) - listed = store.list_active_sessions(workflow_run_id="wfr-1") - assert len(listed) == 1 - assert listed[0].runtime_layer_specs == [] - - -def test_mark_cleaned_sets_status_and_cleaned_at_with_backend_run_id(): - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(), - runtime_layer_specs=_specs(), + context = MagicMock() + session = context.__enter__.return_value + create = MagicMock(return_value=_binding()) + store = WorkflowAgentWorkspaceStore() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, ) - store.mark_cleaned(scope=_scope(), backend_run_id="cleanup-1") + monkeypatch.setattr(store, "_load_execution", MagicMock(return_value=execution)) + monkeypatch.setattr(AgentWorkspaceService, "create_binding", create) - with session_factory.create_session() as session: - row = session.query(WorkflowAgentRuntimeSession).one() - assert row.status == WorkflowAgentRuntimeSessionStatus.CLEANED - assert row.cleaned_at is not None - assert row.backend_run_id == "cleanup-1" + stored = store.load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + assert stored.binding_id == "binding-1" + assert stored.workspace_id == "workspace-1" + assert stored.backend_binding_ref == "backend-binding-1" + assert execution.agent_workspace_binding_id == "binding-1" + assert execution.process_data_dict == { + "existing": "value", + "workflow_agent_binding_id": "workflow-binding-1", + } + assert "agent_workspace_binding_id" not in execution.process_data_dict + assert create.call_args.kwargs["session"] is session + session.commit.assert_called_once_with() -def test_mark_cleaned_preserves_existing_backend_run_id_when_none_given(): - """``backend_run_id=None`` means "leave the previous one in place".""" - store = WorkflowAgentRuntimeSessionStore() - store.save_active_snapshot( - scope=_scope(), - backend_run_id="run-1", - snapshot=_snapshot(), - runtime_layer_specs=_specs(), + get_active = MagicMock(return_value=_binding()) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active) + session.scalar.return_value = execution + resolved = WorkflowAgentSandboxService._resolve_binding( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + session=session, ) - store.mark_cleaned(scope=_scope(), backend_run_id=None) - with session_factory.create_session() as session: - row = session.query(WorkflowAgentRuntimeSession).one() - assert row.status == WorkflowAgentRuntimeSessionStatus.CLEANED - assert row.backend_run_id == "run-1" + assert resolved.id == "binding-1" + owner_scope = get_active.call_args.kwargs["expected_owner_scope"] + assert owner_scope.owner_scope_key == "node-1:workflow-binding-1" -def test_mark_cleaned_is_a_noop_when_no_active_row(): - """No matching ACTIVE row → no-op (already-cleaned rows are not re-touched).""" - store = WorkflowAgentRuntimeSessionStore() - store.mark_cleaned(scope=_scope(), backend_run_id="cleanup-1") - with session_factory.create_session() as session: - assert session.query(WorkflowAgentRuntimeSession).count() == 0 +def test_load_existing_pointer_rejects_missing_workflow_identity(monkeypatch: pytest.MonkeyPatch) -> None: + execution = SimpleNamespace( + agent_workspace_binding_id="binding-1", + process_data=json.dumps({"existing": "value"}), + process_data_dict={"existing": "value"}, + ) + context = MagicMock() + session = context.__enter__.return_value + store = WorkflowAgentWorkspaceStore() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr(store, "_load_execution", MagicMock(return_value=execution)) + get_active = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active) + + with pytest.raises(AgentWorkspaceNotFoundError, match="caller identity is missing"): + store.load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + assert json.loads(execution.process_data) == {"existing": "value"} + get_active.assert_not_called() + session.commit.assert_not_called() -def test_mark_cleaned_is_a_noop_when_workflow_run_id_missing(): - """Without a workflow_run_id we cannot key the row; ignore the call.""" - store = WorkflowAgentRuntimeSessionStore() - store.mark_cleaned(scope=_scope(workflow_run_id=None), backend_run_id="cleanup-1") - # Sanity — no rows created or touched. - with session_factory.create_session() as session: - assert session.query(WorkflowAgentRuntimeSession).count() == 0 +def test_load_existing_pointer_reuses_matching_workflow_identity(monkeypatch: pytest.MonkeyPatch) -> None: + original_process_data = json.dumps( + { + "existing": "value", + "workflow_agent_binding_id": "workflow-binding-1", + } + ) + execution = SimpleNamespace( + agent_workspace_binding_id="binding-1", + process_data=original_process_data, + process_data_dict=json.loads(original_process_data), + ) + context = MagicMock() + session = context.__enter__.return_value + store = WorkflowAgentWorkspaceStore() + create = MagicMock() + binding = _binding() + validate_generation = MagicMock() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr(store, "_load_execution", MagicMock(return_value=execution)) + monkeypatch.setattr(AgentWorkspaceService, "create_binding", create) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) + monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) + + stored = store.load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + assert stored.binding_id == "binding-1" + assert execution.process_data == original_process_data + assert "agent_workspace_binding_id" not in execution.process_data_dict + create.assert_not_called() + session.commit.assert_not_called() + validate_generation.assert_called_once_with( + binding, + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + + +def test_load_existing_pointer_rejects_conflicting_workflow_identity(monkeypatch: pytest.MonkeyPatch) -> None: + execution = SimpleNamespace( + agent_workspace_binding_id="binding-1", + process_data=json.dumps({"workflow_agent_binding_id": "workflow-binding-other"}), + process_data_dict={"workflow_agent_binding_id": "workflow-binding-other"}, + ) + context = MagicMock() + session = context.__enter__.return_value + store = WorkflowAgentWorkspaceStore() + get_active = MagicMock() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr(store, "_load_execution", MagicMock(return_value=execution)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active) + + with pytest.raises(AgentWorkspaceNotFoundError, match="caller identity does not match"): + store.load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + get_active.assert_not_called() + session.commit.assert_not_called() + + +def test_load_or_create_fails_before_binding_create_when_caller_row_is_missing(monkeypatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + create = MagicMock() + sleep = MagicMock() + store = WorkflowAgentWorkspaceStore() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr("core.workflow.nodes.agent_v2.session_store.time.sleep", sleep) + monkeypatch.setattr(session, "scalar", MagicMock(return_value=None)) + monkeypatch.setattr(AgentWorkspaceService, "create_binding", create) + + with pytest.raises(AgentWorkspaceNotFoundError, match="Workflow node execution caller is unavailable"): + store.load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1") + + assert session.scalar.call_count == 60 + assert sleep.call_count == 59 + create.assert_not_called() + session.commit.assert_not_called() + + +def test_load_existing_scope_waits_for_caller_row_to_become_visible(monkeypatch: pytest.MonkeyPatch) -> None: + execution = SimpleNamespace(agent_workspace_binding_id=None) + context = MagicMock() + session = context.__enter__.return_value + session.scalar.side_effect = [None, None, execution] + sleep = MagicMock() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr("core.workflow.nodes.agent_v2.session_store.time.sleep", sleep) + + scope = WorkflowAgentWorkspaceStore().load_existing_node_execution_scope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + ) + + assert scope is None + assert session.scalar.call_count == 3 + assert sleep.call_count == 2 + + +def test_save_snapshot_targets_binding(monkeypatch) -> None: + save = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "save_binding_session_snapshot", save) + snapshot = CompositorSessionSnapshot(layers=[]) + + WorkflowAgentWorkspaceStore().save_active_snapshot(scope=_scope(), binding_id="binding-1", snapshot=snapshot) + + assert save.call_args.kwargs["binding_id"] == "binding-1" + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_retire_workflow_run_only_retires_matching_tenant_and_app( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + matching = _workspace_row() + other_tenant = _workspace_row(workspace_id="workspace-other-tenant", tenant_id="tenant-2") + other_app = _workspace_row( + workspace_id="workspace-other-app", + app_id="app-2", + owner_scope_key="node-2:workflow-binding-2", + ) + sqlite_session.add_all([matching, other_tenant, other_app]) + sqlite_session.commit() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + workspace_ids = WorkflowAgentWorkspaceStore().retire_workflow_run( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + ) + + assert matching.status is AgentWorkingResourceStatus.RETIRED + assert other_tenant.status is AgentWorkingResourceStatus.ACTIVE + assert other_app.status is AgentWorkingResourceStatus.ACTIVE + assert workspace_ids == [matching.id] + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_retire_workflow_run_transitions_active_workspace( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace_row() + binding = _binding_row() + sqlite_session.add_all([workspace, binding]) + sqlite_session.commit() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + workspace_ids = WorkflowAgentWorkspaceStore().retire_workflow_run( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + ) + + assert workspace.status is AgentWorkingResourceStatus.RETIRED + assert binding.status is AgentWorkingResourceStatus.RETIRED + assert workspace_ids == [workspace.id] + + +def test_retire_workflow_run_returns_existing_retired_workspace(monkeypatch) -> None: + workspace = _workspace_row(status=AgentWorkingResourceStatus.RETIRED) + context = MagicMock() + session = context.__enter__.return_value + session.scalars.return_value.all.return_value = [workspace] + retire = MagicMock() + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr(AgentWorkspaceService, "retire_workspace", retire) + + workspace_ids = WorkflowAgentWorkspaceStore().retire_workflow_run( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + ) + + retire.assert_not_called() + assert workspace_ids == [workspace.id] diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_workspace_retirement_layer.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_workspace_retirement_layer.py new file mode 100644 index 00000000000..b5298c23235 --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_workspace_retirement_layer.py @@ -0,0 +1,75 @@ +from unittest.mock import MagicMock + +import pytest + +from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom +from core.workflow.nodes.agent_v2 import workspace_retirement_layer as layer_module +from core.workflow.nodes.agent_v2.workspace_retirement_layer import WorkflowAgentWorkspaceRetirementLayer +from graphon.graph_events import GraphRunSucceededEvent, NodeRunStartedEvent + + +def _run_context() -> DifyRunContext: + return DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="account-1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ) + + +def test_terminal_event_retires_workflow_workspace(monkeypatch: pytest.MonkeyPatch) -> None: + store = MagicMock() + events: list[str] = [] + store.retire_workflow_run.side_effect = lambda **_kwargs: events.append("retire") or ["workspace-1"] + enqueue = MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")) + monkeypatch.setattr(layer_module, "WorkflowAgentWorkspaceStore", MagicMock(return_value=store)) + monkeypatch.setattr(layer_module, "enqueue_agent_resource_collection", enqueue) + layer = WorkflowAgentWorkspaceRetirementLayer(dify_run_context=_run_context()) + layer.initialize(MagicMock(), MagicMock()) + monkeypatch.setattr(layer_module, "get_system_text", lambda *_: "workflow-run-1") + + layer.on_event(MagicMock(spec=GraphRunSucceededEvent)) + + store.retire_workflow_run.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="workflow-run-1", + ) + enqueue.assert_called_once_with(tenant_id="tenant-1", workspace_ids=["workspace-1"]) + assert events == ["retire", "enqueue"] + + +def test_non_terminal_event_does_not_retire_workspace(monkeypatch: pytest.MonkeyPatch) -> None: + store = MagicMock() + monkeypatch.setattr(layer_module, "WorkflowAgentWorkspaceStore", MagicMock(return_value=store)) + layer = WorkflowAgentWorkspaceRetirementLayer(dify_run_context=_run_context()) + + layer.on_event(MagicMock(spec=NodeRunStartedEvent)) + + store.retire_workflow_run.assert_not_called() + + +def test_terminal_retirement_failure_does_not_replace_terminal_event(monkeypatch: pytest.MonkeyPatch) -> None: + store = MagicMock() + store.retire_workflow_run.side_effect = RuntimeError("database unavailable") + log_exception = MagicMock() + enqueue = MagicMock() + monkeypatch.setattr(layer_module, "WorkflowAgentWorkspaceStore", MagicMock(return_value=store)) + monkeypatch.setattr(layer_module.logger, "exception", log_exception) + monkeypatch.setattr(layer_module, "enqueue_agent_resource_collection", enqueue) + layer = WorkflowAgentWorkspaceRetirementLayer(dify_run_context=_run_context()) + layer.initialize(MagicMock(), MagicMock()) + monkeypatch.setattr(layer_module, "get_system_text", lambda *_: "workflow-run-1") + + layer.on_event(MagicMock(spec=GraphRunSucceededEvent)) + + log_exception.assert_called_once_with( + "Failed to retire Workflow Agent Workspaces", + extra={ + "tenant_id": "tenant-1", + "app_id": "app-1", + "workflow_run_id": "workflow-run-1", + }, + ) + enqueue.assert_not_called() diff --git a/api/tests/unit_tests/events/test_app_event_signals.py b/api/tests/unit_tests/events/test_app_event_signals.py index 78e441f7eee..6feedf01273 100644 --- a/api/tests/unit_tests/events/test_app_event_signals.py +++ b/api/tests/unit_tests/events/test_app_event_signals.py @@ -10,6 +10,7 @@ from sqlalchemy.orm import Session from events.app_event import app_was_deleted, app_was_updated from models.account import Account +from models.agent import Agent, AgentWorkspace from models.dataset import AppDatasetJoin from models.model import App, AppMode, AppModelConfig, IconType, InstalledApp from services.app_service import AppService @@ -61,7 +62,7 @@ def _make_collector(target: list[App]): return handler -@pytest.mark.parametrize("sqlite_session", [(App, Account)], indirect=True) +@pytest.mark.parametrize("sqlite_session", [(App, Account, Agent, AgentWorkspace)], indirect=True) @pytest.mark.usefixtures("_mock_deps") class TestAppWasDeletedSignal: def test_sends_signal(self, app_model: App, sqlite_session: Session) -> None: diff --git a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py index 913e8ea6722..4d760748839 100644 --- a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py @@ -1,4 +1,5 @@ -from unittest.mock import patch +import json +from unittest.mock import MagicMock, patch from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( LogstoreAPIWorkflowNodeExecutionRepository, @@ -13,3 +14,28 @@ def test_load_full_process_data_returns_logstore_mapping() -> None: execution.process_data = '{"__dify_retry_history": [{"retry_index": 1}]}' assert repository.load_full_process_data(execution) == {"__dify_retry_history": [{"retry_index": 1}]} + + +def test_get_execution_by_id_keeps_process_data_from_highest_failed_log_version() -> None: + with patch("extensions.logstore.repositories.logstore_api_workflow_node_execution_repository.AliyunLogStore"): + repository = LogstoreAPIWorkflowNodeExecutionRepository(session_maker=None) + repository.logstore_client = MagicMock(supports_pg_protocol=False) + repository.logstore_client.get_logs.return_value = [ + { + "id": "execution-1", + "log_version": "1", + "process_data": "{}", + }, + { + "id": "execution-1", + "log_version": "2", + "status": "failed", + "process_data": json.dumps({"workflow_agent_binding_id": "binding-1"}), + }, + ] + + execution = repository.get_execution_by_id("execution-1") + + assert execution is not None + assert execution.status.value == "failed" + assert execution.process_data_dict == {"workflow_agent_binding_id": "binding-1"} diff --git a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py new file mode 100644 index 00000000000..8ddba55f6ce --- /dev/null +++ b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock, patch + +import pytest + +from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( + LogstoreWorkflowNodeExecutionRepository, +) +from models.account import Account +from models.workflow import WorkflowNodeExecutionTriggeredFrom + + +def test_save_synchronously_writes_sql_when_dual_write_is_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("LOGSTORE_DUAL_WRITE_ENABLED", raising=False) + with ( + patch("extensions.logstore.repositories.logstore_workflow_node_execution_repository.AliyunLogStore"), + patch( + "extensions.logstore.repositories.logstore_workflow_node_execution_repository." + "SQLAlchemyWorkflowNodeExecutionRepository" + ) as sql_repository_type, + ): + repository = LogstoreWorkflowNodeExecutionRepository( + session_factory=MagicMock(), + tenant_id="tenant-1", + user=cast(Account, SimpleNamespace(id="account-1")), + app_id="app-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + execution = MagicMock() + repository.save_synchronously(execution) + + assert repository._enable_dual_write is False + sql_repository_type.return_value.save_synchronously.assert_called_once_with(execution) diff --git a/api/tests/unit_tests/extensions/test_celery_ssl.py b/api/tests/unit_tests/extensions/test_celery_ssl.py index 366e45d86d8..784f6b6417f 100644 --- a/api/tests/unit_tests/extensions/test_celery_ssl.py +++ b/api/tests/unit_tests/extensions/test_celery_ssl.py @@ -193,7 +193,7 @@ class TestCelerySSLConfiguration: assert "redis_backend_use_ssl" in celery_app.conf assert celery_app.conf["redis_backend_use_ssl"] is not None - def test_celery_init_applies_global_keyprefix_to_broker_and_backend_transport(self): + def test_celery_init_applies_global_keyprefix_and_registers_agent_resource_collector(self): mock_config = MagicMock() mock_config.BROKER_USE_SSL = False mock_config.REDIS_KEY_PREFIX = "enterprise-a" @@ -238,3 +238,4 @@ class TestCelerySSLConfiguration: assert celery_app.conf["broker_transport_options"]["global_keyprefix"] == "enterprise-a:" assert celery_app.conf["result_backend_transport_options"]["global_keyprefix"] == "enterprise-a:" + assert "tasks.collect_agent_resources_task" in celery_app.conf["imports"] diff --git a/api/tests/unit_tests/migrations/test_agent_workspace_migration.py b/api/tests/unit_tests/migrations/test_agent_workspace_migration.py new file mode 100644 index 00000000000..7cb92a4d3dc --- /dev/null +++ b/api/tests/unit_tests/migrations/test_agent_workspace_migration.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import importlib.util +from collections.abc import Callable +from pathlib import Path +from typing import Protocol, cast + +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations + +_MIGRATION_PATH = ( + Path(__file__).resolve().parents[3] + / "migrations/versions/2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py" +) + + +class _MigrationModule(Protocol): + op: Operations + upgrade: Callable[[], None] + + +def _load_migration_module() -> _MigrationModule: + spec = importlib.util.spec_from_file_location("agent_workspace_migration", _MIGRATION_PATH) + if spec is None or spec.loader is None: + raise RuntimeError("failed to load Agent Workspace migration") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return cast(_MigrationModule, module) + + +def _create_pre_upgrade_schema(engine: sa.Engine) -> None: + metadata = sa.MetaData() + runtime_sessions = sa.Table( + "agent_runtime_sessions", + metadata, + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("tenant_id", sa.String(36), nullable=False), + sa.Column("conversation_id", sa.String(36)), + sa.Column("workflow_run_id", sa.String(36)), + sa.Column("node_id", sa.String(255)), + sa.Column("binding_id", sa.String(36)), + sa.Column("agent_id", sa.String(36), nullable=False), + sa.Column("agent_config_snapshot_id", sa.String(36)), + sa.Column("home_snapshot_id", sa.String(36), nullable=False), + sa.Column("backend_run_id", sa.String(255)), + sa.Column("status", sa.String(32), nullable=False), + ) + sa.Index("agent_runtime_session_backend_run_idx", runtime_sessions.c.backend_run_id) + sa.Index( + "agent_runtime_session_conversation_lookup_idx", + runtime_sessions.c.tenant_id, + runtime_sessions.c.conversation_id, + runtime_sessions.c.status, + ) + sa.Index( + "agent_runtime_session_conversation_scope_unique", + runtime_sessions.c.tenant_id, + runtime_sessions.c.conversation_id, + runtime_sessions.c.agent_id, + runtime_sessions.c.agent_config_snapshot_id, + runtime_sessions.c.home_snapshot_id, + unique=True, + ) + sa.Index( + "agent_runtime_session_workflow_lookup_idx", + runtime_sessions.c.tenant_id, + runtime_sessions.c.workflow_run_id, + runtime_sessions.c.node_id, + runtime_sessions.c.status, + ) + sa.Index( + "agent_runtime_session_workflow_scope_unique", + runtime_sessions.c.tenant_id, + runtime_sessions.c.workflow_run_id, + runtime_sessions.c.node_id, + runtime_sessions.c.binding_id, + runtime_sessions.c.agent_id, + unique=True, + ) + sa.Table( + "agent_home_snapshots", + metadata, + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("tenant_id", sa.String(36), nullable=False), + sa.Column("agent_id", sa.String(36), nullable=False), + sa.Column("snapshot_ref", sa.String(255), nullable=False), + ) + for table_name in ("conversations", "agent_config_drafts", "workflow_node_executions"): + sa.Table(table_name, metadata, sa.Column("id", sa.String(36), primary_key=True)) + metadata.create_all(engine) + + +def _run_upgrade(module: _MigrationModule, engine: sa.Engine) -> None: + with engine.begin() as connection: + operations = Operations(MigrationContext.configure(connection)) + original_op = module.op + module.op = operations + try: + module.upgrade() + finally: + module.op = original_op + + +def test_upgrade_replaces_runtime_sessions_with_workspace_schema() -> None: + engine = sa.create_engine("sqlite:///:memory:") + _create_pre_upgrade_schema(engine) + + _run_upgrade(_load_migration_module(), engine) + + inspector = sa.inspect(engine) + assert "agent_runtime_sessions" not in inspector.get_table_names() + binding_columns = {column["name"] for column in inspector.get_columns("agent_workspace_bindings")} + assert { + "workspace_id", + "agent_id", + "base_home_snapshot_id", + "agent_config_version_id", + "agent_config_version_kind", + "backend_binding_ref", + "session_snapshot", + "status", + "retired_at", + "pending_form_id", + "pending_tool_call_id", + }.issubset(binding_columns) + assert "active_guard" not in binding_columns + binding_indexes = {index["name"] for index in inspector.get_indexes("agent_workspace_bindings")} + assert "agent_workspace_binding_agent_active_unique" not in binding_indexes + workspace_columns = {column["name"] for column in inspector.get_columns("agent_workspaces")} + assert {"owner_type", "owner_id", "owner_scope_key", "backend_workspace_ref", "status"}.issubset(workspace_columns) + workspace_indexes = {index["name"] for index in inspector.get_indexes("agent_workspaces")} + assert "agent_workspace_owner_active_unique" in workspace_indexes + home_columns = {column["name"] for column in inspector.get_columns("agent_home_snapshots")} + assert {"status", "retired_at"}.issubset(home_columns) + home_indexes = {index["name"] for index in inspector.get_indexes("agent_home_snapshots")} + assert "agent_home_snapshot_status_retired_idx" in home_indexes + for table_name in ("conversations", "agent_config_drafts", "workflow_node_executions"): + caller_columns = {column["name"] for column in inspector.get_columns(table_name)} + assert "agent_workspace_binding_id" in caller_columns diff --git a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py b/api/tests/unit_tests/services/agent/test_agent_dsl_service.py index 4ceecc12a65..1dc8b99eb4a 100644 --- a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py +++ b/api/tests/unit_tests/services/agent/test_agent_dsl_service.py @@ -51,6 +51,7 @@ def _snapshot(*, snapshot_id: str = "snapshot-1", soul: AgentSoulConfig | None = tenant_id="tenant-1", agent_id="agent-1", version=1, + home_snapshot_id="home-1", config_snapshot=soul or AgentSoulConfig(), created_by="account-1", ) @@ -385,7 +386,11 @@ def test_import_workflow_packages_materializes_every_package_binding_as_inline() } for node in graph["nodes"][:3]: node["data"][AGENT_NODE_JOB_DSL_KEY] = {"workflow_prompt": node["id"]} - old_binding = SimpleNamespace(id="old-binding") + old_binding = SimpleNamespace( + id="old-binding", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="old-inline-agent", + ) session = Mock() session.scalars.return_value.all.return_value = [old_binding] service = AgentDslService(session) @@ -406,7 +411,7 @@ def test_import_workflow_packages_materializes_every_package_binding_as_inline() graph="{}", ) - result, warnings = service.import_workflow_packages( + result, warnings, retirement_candidates = service.import_workflow_packages( workflow=workflow, portable_graph=graph, raw_packages={"agent_1": package.model_dump(mode="json")}, @@ -414,6 +419,7 @@ def test_import_workflow_packages_materializes_every_package_binding_as_inline() ) session.delete.assert_called_once_with(old_binding) + assert retirement_candidates == {"old-inline-agent"} assert service._create_imported_inline_agent.call_count == 3 assert [call.kwargs["node_id"] for call in service._create_imported_inline_agent.call_args_list] == [ "roster-1", @@ -678,10 +684,12 @@ def test_resolve_package_soul_preserves_existing_and_marks_missing_knowledge(mon } -def test_create_snapshot_increments_version_and_records_revision() -> None: +def test_create_snapshot_increments_version_and_records_revision(monkeypatch: pytest.MonkeyPatch) -> None: session = Mock() session.scalar.return_value = 2 service = AgentDslService(session) + create_initial = Mock(return_value=SimpleNamespace(id="home-3")) + monkeypatch.setattr("services.agent.dsl_service.AgentHomeSnapshotService.create_initial", create_initial) snapshot = service._create_snapshot( tenant_id="tenant-1", @@ -692,6 +700,12 @@ def test_create_snapshot_increments_version_and_records_revision() -> None: ) assert snapshot.version == 3 + assert snapshot.home_snapshot_id == "home-3" + create_initial.assert_called_once_with( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + ) assert isinstance(session.add.call_args_list[0].args[0], AgentConfigSnapshot) revision = session.add.call_args_list[1].args[0] assert isinstance(revision, AgentConfigRevision) diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py index f881e8770dd..4c4d3717d6c 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -1,26 +1,30 @@ import json +from contextlib import nullcontext from datetime import UTC, datetime from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest -from agenton.compositor import CompositorSessionSnapshot -from dify_agent.protocol import RuntimeLayerSpec from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidationError from models.agent import ( Agent, AgentConfigDraft, AgentConfigDraftType, + AgentConfigRevision, AgentConfigRevisionOperation, AgentConfigSnapshot, + AgentConfigVersionKind, AgentDebugConversation, AgentDriveFile, + AgentHomeSnapshot, AgentKind, AgentScope, AgentSource, AgentStatus, + AgentWorkspaceOwnerType, WorkflowAgentBindingType, WorkflowAgentNodeBinding, ) @@ -39,6 +43,7 @@ from services.agent.agent_soul_state import agent_soul_has_model from services.agent.composer_service import AgentComposerService from services.agent.composer_validator import ComposerConfigValidator from services.agent.errors import ( + AgentBuildSandboxNotFoundError, AgentModelNotConfiguredError, AgentNameConflictError, AgentNotFoundError, @@ -46,8 +51,10 @@ from services.agent.errors import ( AgentVersionNotFoundError, InvalidComposerConfigError, ) +from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.roster_service import AgentRosterService from services.agent.workflow_publish_service import WorkflowAgentPublishService +from services.agent.workspace_service import AgentWorkspaceService from services.app_service import AppListParams, AppService from services.entities.agent_entities import AgentSoulConfig, ComposerSavePayload, ComposerSaveStrategy, ComposerVariant @@ -60,6 +67,25 @@ class FakeScalarResult: return self.values +class FakeTransaction: + def __init__(self, session: "FakeSession") -> None: + self._session = session + + def __enter__(self) -> None: + return None + + def __exit__(self, exc_type, exc, traceback) -> bool: + if exc_type is not None: + self._session.rollback() + return False + try: + self._session.commit() + except Exception: + self._session.rollback() + raise + return False + + class FakeSession: def __init__(self, *, scalars=None, scalar=None): self._scalars = list(scalars or []) @@ -69,6 +95,10 @@ class FakeSession: self.commits = 0 self.flushes = 0 self.rollbacks = 0 + self.info = {} + + def begin(self): + return FakeTransaction(self) def scalar(self, _stmt): if self._scalar: @@ -99,6 +129,15 @@ class FakeSession: self.rollbacks += 1 +@pytest.fixture +def _stub_home_snapshot_backend(monkeypatch: pytest.MonkeyPatch): + create_initial = MagicMock(return_value=SimpleNamespace(id="home-initial", snapshot_ref="backend-home-initial")) + create_for_build_apply = MagicMock(return_value=SimpleNamespace(id="home-build", snapshot_ref="backend-home-build")) + monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_for_build_apply) + return SimpleNamespace(create_initial=create_initial, create_for_build_apply=create_for_build_apply) + + def _agent_soul_with_model() -> AgentSoulConfig: return AgentSoulConfig.model_validate( { @@ -309,7 +348,7 @@ def test_workflow_inline_debug_conversation_seed(monkeypatch: pytest.MonkeyPatch def __init__(self, session): captured["session"] = session - def get_or_create_agent_app_debug_conversation_id(self, **kwargs): + def get_or_create_build_conversation(self, **kwargs): captured.update(kwargs) return "debug-conversation-1" @@ -449,6 +488,69 @@ def test_save_workflow_composer_dispatches_save_strategy(monkeypatch, strategy, assert fake_session.flushes >= 1 +def test_save_workflow_composer_commits_before_retiring_replaced_inline_agent(monkeypatch) -> None: + session = FakeSession() + events: list[str] = [] + old_binding = SimpleNamespace( + agent_id="old-inline-agent", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + ) + new_binding = SimpleNamespace( + agent_id="new-agent", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + current_snapshot_id="version-1", + ) + monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **_kwargs: SimpleNamespace(id="workflow-1")) + monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **_kwargs: old_binding) + monkeypatch.setattr(AgentComposerService, "_save_as_new_agent", lambda **_kwargs: new_binding) + monkeypatch.setattr( + AgentComposerService, + "_get_agent_if_present", + lambda **_kwargs: SimpleNamespace(id="new-agent", active_config_snapshot_id="version-1"), + ) + monkeypatch.setattr( + AgentComposerService, + "_get_version_if_present", + lambda **_kwargs: SimpleNamespace(id="version-1"), + ) + monkeypatch.setattr(AgentComposerService, "_serialize_workflow_state", lambda **_kwargs: {"state": "ok"}) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None) + monkeypatch.setattr(AgentComposerService, "collect_validation_findings", lambda **_kwargs: {}) + session.commit = lambda: events.append("commit") # type: ignore[method-assign] + + def retire_unowned(**kwargs): + assert kwargs["agent_ids"] == {"old-inline-agent"} + events.append("retire") + return ["binding-1"], ["home-1"] + + monkeypatch.setattr(composer_service.WorkflowAgentRetirementService, "retire_unowned", retire_unowned) + monkeypatch.setattr( + composer_service, + "enqueue_agent_resource_collection", + MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), + ) + payload = ComposerSavePayload.model_validate( + { + "variant": ComposerVariant.WORKFLOW, + "save_strategy": ComposerSaveStrategy.SAVE_AS_NEW_AGENT, + "agent_soul": _agent_soul_with_model().model_dump(mode="json"), + "new_agent_name": "New Agent", + "soul_lock": {"locked": False}, + } + ) + + AgentComposerService.save_workflow_composer( + session=session, + tenant_id="tenant-1", + app_id="app-1", + node_id="node-1", + account_id="account-1", + payload=payload, + ) + + assert events == ["commit", "retire", "enqueue"] + + def test_save_workflow_composer_rejects_agent_app_variant(): session = FakeSession() payload = ComposerSavePayload.model_validate( @@ -510,9 +612,14 @@ def test_publish_save_strategies_run_publish_validation(strategy: ComposerSaveSt composer_service._validate_composer_payload_for_strategy(_duplicate_env_secret_payload(strategy)) +@pytest.mark.usefixtures("_stub_home_snapshot_backend") def test_save_agent_app_composer_creates_agent_when_missing(monkeypatch: pytest.MonkeyPatch): fake_session = FakeSession(scalar=[None]) - saved_draft = SimpleNamespace(id="draft-1", config_snapshot_dict={"prompt": {"system_prompt": "x"}}) + saved_draft = SimpleNamespace( + id="draft-1", + home_snapshot_id="home-initial", + config_snapshot_dict={"prompt": {"system_prompt": "x"}}, + ) session = fake_session monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_draft_save_payload", lambda payload: None) @@ -538,7 +645,8 @@ def test_save_agent_app_composer_creates_agent_when_missing(monkeypatch: pytest. assert result.pop("validation") == {"warnings": [], "knowledge_retrieval_placeholder": []} assert result == {"loaded": True} assert fake_session.added[0].name == "Analyst" - assert fake_session.added[0].active_config_snapshot_id is None + assert fake_session.added[0].active_config_snapshot_id == fake_session.added[1].id + assert fake_session.added[1].home_snapshot_id == "home-initial" assert fake_session.added[0].active_config_is_published is False assert fake_session.flushes >= 1 @@ -600,7 +708,9 @@ def test_save_agent_app_composer_updates_normal_draft(monkeypatch: pytest.Monkey active_config_is_published=True, updated_by=None, ) - active_version = SimpleNamespace(config_snapshot_dict=AgentSoulConfig().model_dump(mode="json")) + active_version = SimpleNamespace( + home_snapshot_id="home-initial", config_snapshot_dict=AgentSoulConfig().model_dump(mode="json") + ) fake_session = FakeSession(scalar=[agent]) saved = {} @@ -609,7 +719,7 @@ def test_save_agent_app_composer_updates_normal_draft(monkeypatch: pytest.Monkey monkeypatch.setattr( AgentComposerService, "_save_agent_draft", - lambda **kwargs: saved.update(kwargs) or SimpleNamespace(id="draft-1"), + lambda **kwargs: saved.update(kwargs) or SimpleNamespace(id="draft-1", home_snapshot_id="home-initial"), ) monkeypatch.setattr(AgentComposerService, "_get_version_if_present", lambda **_kwargs: active_version) monkeypatch.setattr( @@ -652,7 +762,9 @@ def test_save_agent_app_composer_keeps_published_when_draft_matches_active_snaps active_config_is_published=False, updated_by=None, ) - active_version = SimpleNamespace(config_snapshot_dict=agent_soul.model_dump(mode="json")) + active_version = SimpleNamespace( + home_snapshot_id="home-initial", config_snapshot_dict=agent_soul.model_dump(mode="json") + ) fake_session = FakeSession(scalar=[agent, "publish-revision-1"]) session = fake_session @@ -660,7 +772,7 @@ def test_save_agent_app_composer_keeps_published_when_draft_matches_active_snaps monkeypatch.setattr( AgentComposerService, "_save_agent_draft", - lambda **_kwargs: SimpleNamespace(id="draft-1"), + lambda **_kwargs: SimpleNamespace(id="draft-1", home_snapshot_id="home-initial"), ) monkeypatch.setattr(AgentComposerService, "_get_version_if_present", lambda **_kwargs: active_version) monkeypatch.setattr( @@ -757,19 +869,26 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest. draft_type=AgentConfigDraftType.DRAFT, draft_owner_key="", base_snapshot_id="version-1", + home_snapshot_id="home-1", config_snapshot=_agent_soul_with_model(), ) version = SimpleNamespace(id="version-2") fake_session = FakeSession(scalar=[agent, draft]) created: dict[str, object] = {} + calls: list[str] = [] session = fake_session monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda payload: None) monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **kwargs: None) + monkeypatch.setattr( + composer_service, + "validate_home_snapshot_binding", + lambda **kwargs: calls.append("validate_home"), + ) monkeypatch.setattr( AgentComposerService, "_create_config_version", - lambda **kwargs: created.update(kwargs) or version, + lambda **kwargs: calls.append("create_version") or created.update(kwargs) or version, ) monkeypatch.setattr(AgentComposerService, "_serialize_version", lambda _version: {"id": _version.id}) @@ -786,12 +905,76 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest. assert result["draft"]["base_snapshot_id"] == "version-2" assert created["operation"] == AgentConfigRevisionOperation.PUBLISH_DRAFT assert created["previous_snapshot_id"] == "version-1" + assert created["home_snapshot_id"] == "home-1" + assert calls == ["validate_home", "create_version"] assert agent.active_config_snapshot_id == "version-2" assert agent.active_config_has_model is True assert agent.active_config_is_published is True assert fake_session.flushes >= 1 +def test_repeated_publish_reuses_normal_draft_home_without_creating_resources( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + active_config_snapshot_id="version-1", + ) + draft = AgentConfigDraft( + id="draft-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-1", + config_snapshot=_agent_soul_with_model(), + ) + session = FakeSession(scalar=[agent, draft, agent, draft]) + published_homes: list[str] = [] + versions = iter([SimpleNamespace(id="version-2"), SimpleNamespace(id="version-3")]) + create_initial = MagicMock() + create_from_build = MagicMock() + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None) + monkeypatch.setattr(composer_service, "validate_home_snapshot_binding", lambda **_kwargs: None) + monkeypatch.setattr( + AgentComposerService, + "_create_config_version", + lambda **kwargs: published_homes.append(kwargs["home_snapshot_id"]) or next(versions), + ) + monkeypatch.setattr(AgentComposerService, "_serialize_version", lambda version: {"id": version.id}) + monkeypatch.setattr(AgentComposerService, "_serialize_draft", lambda value: {"id": value.id}) + monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_from_build) + + first = AgentComposerService.publish_agent_app_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + second = AgentComposerService.publish_agent_app_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + assert first["active_config_snapshot_id"] == "version-2" + assert second["active_config_snapshot_id"] == "version-3" + assert published_homes == ["home-1", "home-1"] + assert draft.home_snapshot_id == "home-1" + create_initial.assert_not_called() + create_from_build.assert_not_called() + + def test_agent_app_build_draft_checkout_and_apply_use_user_isolated_draft(monkeypatch: pytest.MonkeyPatch): agent = Agent( id="agent-1", @@ -812,6 +995,7 @@ def test_agent_app_build_draft_checkout_and_apply_use_user_isolated_draft(monkey account_id=None, draft_owner_key="", base_snapshot_id="version-1", + home_snapshot_id="home-initial", config_snapshot=_agent_soul_with_model(), ) fake_session = FakeSession(scalar=[agent, normal_draft, None]) @@ -830,14 +1014,26 @@ def test_agent_app_build_draft_checkout_and_apply_use_user_isolated_draft(monkey assert checked_out["draft"]["draft_type"] == AgentConfigDraftType.DEBUG_BUILD.value assert checked_out["draft"]["account_id"] == "account-1" assert checked_out["draft"]["base_snapshot_id"] == "version-1" + assert build_draft.home_snapshot_id == "home-initial" assert checked_out["agent_soul"] == normal_draft.config_snapshot_dict assert fake_session.flushes >= 1 + assert fake_session.commits == 1 - active_version = SimpleNamespace(config_snapshot_dict=build_draft.config_snapshot_dict) + active_version = SimpleNamespace( + home_snapshot_id=build_draft.home_snapshot_id, config_snapshot_dict=build_draft.config_snapshot_dict + ) fake_session = FakeSession( scalar=[agent, build_draft, normal_draft, active_version, "publish-revision-1"], ) session = fake_session + source_binding_id = "binding-1" + build_draft.agent_workspace_binding_id = source_binding_id + create_home = MagicMock(return_value=SimpleNamespace(id="home-build", snapshot_ref="backend-home-build")) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_home) + retire_binding = MagicMock() + enqueue_collection = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) applied = AgentComposerService.apply_agent_app_build_draft( session=session, @@ -849,9 +1045,857 @@ def test_agent_app_build_draft_checkout_and_apply_use_user_isolated_draft(monkey assert applied["result"] == "success" assert applied["draft"]["id"] == normal_draft.id assert normal_draft.config_snapshot_dict == build_draft.config_snapshot_dict - assert agent.active_config_is_published is True + assert normal_draft.home_snapshot_id == "home-build" + assert agent.active_config_is_published is False assert fake_session.deleted == [build_draft] assert fake_session.flushes >= 1 + create_home.assert_called_once_with( + session=session, + build_draft=build_draft, + ) + retire_binding.assert_called_once_with(session=session, tenant_id="tenant-1", binding_id=source_binding_id) + enqueue_collection.assert_called_once_with(tenant_id="tenant-1", binding_ids=[source_binding_id]) + + +@pytest.mark.parametrize( + ("scope", "source", "app_id", "backing_app_id", "expected_runtime_app_id"), + [ + (AgentScope.ROSTER, AgentSource.AGENT_APP, "app-1", None, "app-1"), + ( + AgentScope.WORKFLOW_ONLY, + AgentSource.WORKFLOW, + "workflow-app-1", + "runtime-app-1", + "runtime-app-1", + ), + ], +) +def test_force_build_draft_checkout_collects_retired_binding_after_commit( + monkeypatch: pytest.MonkeyPatch, + scope: AgentScope, + source: AgentSource, + app_id: str, + backing_app_id: str | None, + expected_runtime_app_id: str, +) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=scope, + source=source, + status=AgentStatus.ACTIVE, + app_id=app_id, + backing_app_id=backing_app_id, + ) + normal_draft = AgentConfigDraft( + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-1", + config_snapshot=_agent_soul_with_model(), + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-1", + agent_workspace_binding_id="binding-1", + config_snapshot=_agent_soul_with_model(), + ) + binding = SimpleNamespace( + agent_id=agent.id, + base_home_snapshot_id="home-1", + agent_config_version_id=build_draft.id, + agent_config_version_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) + get_active_binding = MagicMock(return_value=binding) + retire_binding = MagicMock(return_value="binding-1") + enqueue_collection = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + reuse_session = FakeSession(scalar=[agent, normal_draft, build_draft]) + AgentComposerService.checkout_agent_app_build_draft( + session=reuse_session, + tenant_id="tenant-1", + agent_id=agent.id, + account_id="account-1", + ) + + assert reuse_session.commits == 1 + retire_binding.assert_not_called() + enqueue_collection.assert_not_called() + + session = FakeSession(scalar=[agent, normal_draft, build_draft]) + AgentComposerService.checkout_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id=agent.id, + account_id="account-1", + force=True, + ) + + assert build_draft.agent_workspace_binding_id is None + assert session.commits == 1 + owner_scope = get_active_binding.call_args.kwargs["expected_owner_scope"] + assert owner_scope.app_id == expected_runtime_app_id + assert owner_scope.owner_type is AgentWorkspaceOwnerType.BUILD_DRAFT + assert owner_scope.owner_id == build_draft.id + retire_binding.assert_called_once_with( + session=session, + tenant_id="tenant-1", + binding_id="binding-1", + ) + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=("binding-1",), + ) + + +def test_force_build_draft_checkout_rejects_unavailable_pointed_binding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app-1", + ) + normal_draft = AgentConfigDraft( + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-1", + config_snapshot=_agent_soul_with_model(), + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-1", + agent_workspace_binding_id="binding-missing", + config_snapshot=_agent_soul_with_model(), + ) + session = FakeSession(scalar=[agent, normal_draft, build_draft]) + retire_binding = MagicMock() + enqueue_collection = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None)) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(AgentBuildSandboxNotFoundError): + AgentComposerService.checkout_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id=agent.id, + account_id="account-1", + force=True, + ) + + assert build_draft.agent_workspace_binding_id == "binding-missing" + assert session.commits == 0 + assert session.rollbacks == 1 + retire_binding.assert_not_called() + enqueue_collection.assert_not_called() + + +def test_build_apply_checkpoints_binding_updates_normal_draft_then_collects(monkeypatch: pytest.MonkeyPatch) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-old", + agent_workspace_binding_id="binding-1", + config_snapshot=AgentSoulConfig(), + ) + normal_draft = AgentConfigDraft( + id="draft-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-old", + config_snapshot=AgentSoulConfig(), + ) + source_binding = SimpleNamespace( + id="binding-1", + backend_binding_ref="backend-binding-1", + agent_id="agent-1", + base_home_snapshot_id="home-old", + agent_config_version_id="build-1", + agent_config_version_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) + session = FakeSession( + scalar=[agent, build_draft, SimpleNamespace(app_id="app-1", backing_app_id=None), source_binding] + ) + client = MagicMock() + client.create_home_snapshot_from_binding_sync.return_value = SimpleNamespace(snapshot_ref="snapshot-ref-new") + lifecycle: list[str] = [] + original_commit = session.commit + + def commit() -> None: + original_commit() + lifecycle.append("commit") + + session.commit = commit # type: ignore[method-assign] + retire = MagicMock(return_value=source_binding.id) + enqueue_collection = MagicMock(side_effect=lambda **_kwargs: lifecycle.append("enqueue")) + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None) + monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client)) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", lambda **_kwargs: normal_draft) + monkeypatch.setattr(AgentComposerService, "_agent_soul_matches_active_config", lambda **_kwargs: False) + monkeypatch.setattr(AgentComposerService, "_serialize_draft", lambda draft: {"id": draft.id}) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + result = AgentComposerService.apply_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + request = client.create_home_snapshot_from_binding_sync.call_args.args[0] + assert request.backend_binding_ref == "backend-binding-1" + created_home = next(row for row in session.added if isinstance(row, AgentHomeSnapshot)) + assert created_home.snapshot_ref == "snapshot-ref-new" + assert normal_draft.home_snapshot_id == created_home.id + retire.assert_called_once_with(session=session, tenant_id="tenant-1", binding_id="binding-1") + enqueue_collection.assert_called_once_with(tenant_id="tenant-1", binding_ids=["binding-1"]) + assert lifecycle == ["commit", "enqueue"] + assert result == {"result": "success", "draft": {"id": "draft-1"}} + + +def test_build_apply_retires_normal_preview_binding_before_replacing_draft_home( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app-1", + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-build", + agent_workspace_binding_id="binding-build", + config_snapshot=AgentSoulConfig(), + ) + normal_draft = AgentConfigDraft( + id="draft-1", + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-preview-old", + config_snapshot=AgentSoulConfig(), + ) + preview_mapping = SimpleNamespace( + app_id="app-1", + account_id="account-2", + conversation_id="conversation-preview", + ) + empty_preview_mapping = SimpleNamespace( + app_id="app-1", + account_id="account-3", + conversation_id="conversation-preview-empty", + ) + preview_conversation = SimpleNamespace( + id="conversation-preview", + agent_workspace_binding_id="binding-preview", + ) + empty_preview_conversation = SimpleNamespace( + id="conversation-preview-empty", + agent_workspace_binding_id=None, + ) + preview_binding = SimpleNamespace( + id="binding-preview", + agent_id=agent.id, + base_home_snapshot_id="home-preview-old", + agent_config_version_id=normal_draft.id, + agent_config_version_kind=AgentConfigVersionKind.DRAFT, + ) + session = FakeSession( + scalar=[agent, build_draft, preview_conversation, empty_preview_conversation], + scalars=[[preview_mapping, empty_preview_mapping]], + ) + lifecycle: list[str] = [] + original_commit = session.commit + + def commit() -> None: + original_commit() + lifecycle.append("commit") + + session.commit = commit # type: ignore[method-assign] + monkeypatch.setattr( + AgentHomeSnapshotService, + "create_for_build_apply", + MagicMock(return_value=SimpleNamespace(id="home-applied")), + ) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", MagicMock(return_value=normal_draft)) + monkeypatch.setattr(AgentComposerService, "_agent_soul_matches_active_config", MagicMock(return_value=False)) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", MagicMock()) + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", MagicMock()) + get_active_binding = MagicMock(return_value=preview_binding) + retire_binding = MagicMock(side_effect=["binding-preview", "binding-build"]) + validate_generation = MagicMock() + enqueue_collection = MagicMock(side_effect=lambda **_kwargs: lifecycle.append("enqueue")) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) + monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + AgentComposerService.apply_agent_app_build_draft( + session=session, + tenant_id=agent.tenant_id, + agent_id=agent.id, + account_id="account-1", + ) + + assert preview_conversation.agent_workspace_binding_id is None + validate_generation.assert_called_once_with( + preview_binding, + base_home_snapshot_id="home-preview-old", + agent_config_version_id=normal_draft.id, + agent_config_version_kind=AgentConfigVersionKind.DRAFT, + ) + enqueue_collection.assert_called_once_with( + tenant_id=agent.tenant_id, + binding_ids=["binding-preview", "binding-build"], + ) + assert get_active_binding.call_count == 1 + assert get_active_binding.call_args.kwargs["binding_id"] == "binding-preview" + assert retire_binding.call_args_list == [ + call(session=session, tenant_id=agent.tenant_id, binding_id="binding-preview"), + call(session=session, tenant_id=agent.tenant_id, binding_id="binding-build"), + ] + assert lifecycle == ["commit", "enqueue"] + + +def test_build_apply_validates_before_resolving_or_snapshotting_sandbox(monkeypatch: pytest.MonkeyPatch) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-1", + config_snapshot=_agent_soul_with_model(), + ) + validation = MagicMock(side_effect=InvalidComposerConfigError("invalid Build Draft")) + create_home = MagicMock() + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", validation) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_home) + + with pytest.raises(InvalidComposerConfigError, match="invalid Build Draft"): + AgentComposerService.apply_agent_app_build_draft( + session=FakeSession(scalar=[agent, build_draft]), + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + validation.assert_called_once() + create_home.assert_not_called() + + +def test_build_apply_without_model_snapshots_source_sandbox_and_updates_normal_draft( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-old", + agent_workspace_binding_id="binding-1", + config_snapshot=AgentSoulConfig(), + ) + normal_draft = AgentConfigDraft( + id="draft-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-old", + config_snapshot=AgentSoulConfig(), + ) + session = FakeSession(scalar=[agent, build_draft]) + create_home = MagicMock(return_value=SimpleNamespace(id="home-new", snapshot_ref="backend-home-new")) + validate_knowledge = MagicMock() + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", validate_knowledge) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_home) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", lambda **_kwargs: normal_draft) + monkeypatch.setattr(AgentComposerService, "_agent_soul_matches_active_config", lambda **_kwargs: False) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", MagicMock()) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", MagicMock()) + monkeypatch.setattr(AgentComposerService, "_serialize_draft", lambda draft: {"id": draft.id}) + + result = AgentComposerService.apply_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + create_home.assert_called_once_with( + session=session, + build_draft=build_draft, + ) + validate_knowledge.assert_called_once() + assert normal_draft.home_snapshot_id == "home-new" + assert session.deleted == [build_draft] + assert session.commits == 1 + assert result == {"result": "success", "draft": {"id": "draft-1"}} + + +def test_build_apply_requires_retained_sandbox_before_creating_home(monkeypatch: pytest.MonkeyPatch) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-old", + config_snapshot=AgentSoulConfig(), + ) + create_home = MagicMock() + save_draft = MagicMock() + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_home) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", save_draft) + session = FakeSession(scalar=[agent, build_draft]) + + with pytest.raises(AgentBuildSandboxNotFoundError): + AgentComposerService.apply_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + create_home.assert_not_called() + save_draft.assert_not_called() + assert build_draft.home_snapshot_id == "home-old" + assert session.deleted == [] + assert session.commits == 0 + + +def test_build_apply_fails_when_locked_source_cannot_be_retired(monkeypatch: pytest.MonkeyPatch) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-old", + agent_workspace_binding_id="binding-1", + config_snapshot=AgentSoulConfig(), + ) + normal_draft = AgentConfigDraft( + id="draft-1", + tenant_id="tenant-1", + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-old", + config_snapshot=AgentSoulConfig(), + ) + session = FakeSession(scalar=[agent, build_draft]) + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None) + monkeypatch.setattr( + AgentHomeSnapshotService, + "create_for_build_apply", + MagicMock(return_value=SimpleNamespace(id="home-new", snapshot_ref="backend-home-new")), + ) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", MagicMock(return_value=normal_draft)) + monkeypatch.setattr(AgentComposerService, "_agent_soul_matches_active_config", lambda **_kwargs: False) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", MagicMock(return_value=None)) + enqueue_collection = MagicMock() + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(AgentBuildSandboxNotFoundError): + AgentComposerService.apply_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id=agent.id, + account_id="account-1", + ) + + assert session.deleted == [] + assert session.commits == 0 + assert session.rollbacks == 1 + enqueue_collection.assert_not_called() + + +def test_build_apply_home_create_failure_leaves_drafts_untouched(monkeypatch: pytest.MonkeyPatch) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-old", + agent_workspace_binding_id="binding-1", + config_snapshot=AgentSoulConfig(), + ) + normal_draft = AgentConfigDraft( + id="draft-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-old", + config_snapshot=AgentSoulConfig(), + ) + save_draft = MagicMock(return_value=normal_draft) + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None) + monkeypatch.setattr( + AgentHomeSnapshotService, + "create_for_build_apply", + MagicMock(side_effect=RuntimeError("snapshot failed")), + ) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", save_draft) + session = FakeSession(scalar=[agent, build_draft]) + + with pytest.raises(RuntimeError, match="snapshot failed"): + AgentComposerService.apply_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + save_draft.assert_not_called() + assert build_draft.home_snapshot_id == "home-old" + assert normal_draft.home_snapshot_id == "home-old" + assert session.deleted == [] + assert session.commits == 0 + + +def test_build_apply_commit_failure_rolls_back_and_preserves_physical_home( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-old", + agent_workspace_binding_id="binding-1", + config_snapshot=AgentSoulConfig(), + ) + normal_draft = AgentConfigDraft( + id="draft-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + home_snapshot_id="home-old", + config_snapshot=AgentSoulConfig(), + ) + session = FakeSession(scalar=[agent, build_draft]) + + def fail_commit() -> None: + session.commits += 1 + raise RuntimeError("commit failed") + + session.commit = fail_commit # type: ignore[method-assign] + delete_home = MagicMock() + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None) + monkeypatch.setattr( + AgentHomeSnapshotService, + "create_for_build_apply", + lambda **_kwargs: SimpleNamespace(id="home-new", snapshot_ref="backend-home-new"), + ) + monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete_home) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", lambda **_kwargs: normal_draft) + monkeypatch.setattr(AgentComposerService, "_agent_soul_matches_active_config", lambda **_kwargs: False) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", MagicMock(return_value="binding-1")) + enqueue_collection = MagicMock() + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(RuntimeError, match="commit failed"): + AgentComposerService.apply_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + assert session.commits == 1 + assert session.rollbacks == 1 + delete_home.assert_not_called() + enqueue_collection.assert_not_called() + + +def test_build_draft_save_and_discard_do_not_manage_home_resources(monkeypatch: pytest.MonkeyPatch) -> None: + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + ) + build_draft = AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-existing", + config_snapshot=AgentSoulConfig(), + ) + create_initial = MagicMock() + create_from_build = MagicMock() + delete_home = MagicMock() + monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_from_build) + monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete_home) + monkeypatch.setattr(AgentComposerService, "_require_agent", lambda **_kwargs: agent) + monkeypatch.setattr(AgentComposerService, "_save_agent_draft", lambda **_kwargs: build_draft) + monkeypatch.setattr(AgentComposerService, "_serialize_build_draft_state", lambda draft: {"id": draft.id}) + save_session = FakeSession() + + AgentComposerService.save_agent_app_build_draft( + session=save_session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + payload=ComposerSavePayload( + variant=ComposerVariant.AGENT_APP, + save_strategy=ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION, + agent_soul=AgentSoulConfig(), + ), + ) + AgentComposerService.discard_agent_app_build_draft( + session=FakeSession(scalar=[build_draft]), + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + create_initial.assert_not_called() + create_from_build.assert_not_called() + delete_home.assert_not_called() + + +def test_build_discard_retires_then_commits_before_enqueue(monkeypatch: pytest.MonkeyPatch) -> None: + agent = SimpleNamespace(id="agent-1", app_id="app-1", backing_app_id=None) + build_draft = SimpleNamespace( + id="build-1", + agent_id=agent.id, + home_snapshot_id="home-1", + agent_workspace_binding_id="binding-1", + ) + binding = SimpleNamespace( + agent_id=agent.id, + base_home_snapshot_id="home-1", + agent_config_version_id=build_draft.id, + agent_config_version_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) + session = FakeSession(scalar=[agent, build_draft]) + events: list[str] = [] + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) + monkeypatch.setattr( + AgentWorkspaceService, + "retire_binding", + MagicMock(side_effect=lambda **_kwargs: events.append("retire") or "binding-1"), + ) + session.commit = lambda: events.append("commit") # type: ignore[method-assign] + monkeypatch.setattr( + composer_service, + "enqueue_agent_resource_collection", + MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), + ) + + AgentComposerService.discard_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + assert events == ["retire", "commit", "enqueue"] + + +def test_build_discard_commit_failure_does_not_enqueue(monkeypatch: pytest.MonkeyPatch) -> None: + agent = SimpleNamespace(id="agent-1", app_id="app-1", backing_app_id=None) + build_draft = SimpleNamespace( + id="build-1", + agent_id=agent.id, + home_snapshot_id="home-1", + agent_workspace_binding_id="binding-1", + ) + binding = SimpleNamespace( + agent_id=agent.id, + base_home_snapshot_id="home-1", + agent_config_version_id=build_draft.id, + agent_config_version_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) + session = FakeSession(scalar=[agent, build_draft]) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", MagicMock(return_value="binding-1")) + session.commit = MagicMock(side_effect=RuntimeError("commit failed")) # type: ignore[method-assign] + enqueue_collection = MagicMock() + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(RuntimeError, match="commit failed"): + AgentComposerService.discard_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + enqueue_collection.assert_not_called() + + +def test_build_discard_rejects_unavailable_pointed_binding(monkeypatch: pytest.MonkeyPatch) -> None: + agent = SimpleNamespace(id="agent-1", app_id="app-1", backing_app_id=None) + build_draft = SimpleNamespace( + id="build-1", + agent_id=agent.id, + home_snapshot_id="home-1", + agent_workspace_binding_id="binding-missing", + ) + session = FakeSession(scalar=[agent, build_draft]) + retire_binding = MagicMock() + enqueue_collection = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None)) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(AgentBuildSandboxNotFoundError): + AgentComposerService.discard_agent_app_build_draft( + session=session, + tenant_id="tenant-1", + agent_id=agent.id, + account_id="account-1", + ) + + assert build_draft.agent_workspace_binding_id == "binding-missing" + assert session.deleted == [] + assert session.commits == 0 + assert session.rollbacks == 1 + retire_binding.assert_not_called() + enqueue_collection.assert_not_called() @pytest.mark.parametrize( @@ -951,6 +1995,7 @@ def test_agent_app_build_draft_apply_marks_unpublished_when_build_draft_differs( account_id="account-1", draft_owner_key="account-1", base_snapshot_id="version-1", + agent_workspace_binding_id="binding-1", config_snapshot=build_agent_soul, ) normal_draft = AgentConfigDraft( @@ -962,9 +2007,15 @@ def test_agent_app_build_draft_apply_marks_unpublished_when_build_draft_differs( base_snapshot_id="version-1", config_snapshot=active_agent_soul, ) - active_version = SimpleNamespace(config_snapshot_dict=active_agent_soul.model_dump(mode="json")) + active_version = SimpleNamespace( + home_snapshot_id="home-initial", config_snapshot_dict=active_agent_soul.model_dump(mode="json") + ) fake_session = FakeSession(scalar=[agent, build_draft, normal_draft, active_version]) session = fake_session + create_home = MagicMock(return_value=SimpleNamespace(id="home-build", snapshot_ref="backend-home-build")) + monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_home) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", MagicMock()) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", MagicMock()) AgentComposerService.apply_agent_app_build_draft( session=session, @@ -977,6 +2028,10 @@ def test_agent_app_build_draft_apply_marks_unpublished_when_build_draft_differs( assert agent.active_config_is_published is False assert fake_session.deleted == [build_draft] assert fake_session.flushes >= 1 + create_home.assert_called_once_with( + session=session, + build_draft=build_draft, + ) def test_agent_app_composer_candidates_and_impact(monkeypatch: pytest.MonkeyPatch): @@ -1161,6 +2216,7 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk icon_type="emoji", icon="source", icon_background="#FFFFFF", + scope=AgentScope.WORKFLOW_ONLY, ) create_roster_calls = [] copy_drive_calls = [] @@ -1189,6 +2245,7 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk tenant_id="tenant-1", agent_id="roster-agent-1", version=1, + home_snapshot_id="home-source", config_snapshot='{"prompt":{"system_prompt":"old"}}', ), ) @@ -1197,7 +2254,6 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk "_create_config_version", lambda **kwargs: AgentConfigSnapshot(id="new-version-1", version=2), ) - payload = ComposerSavePayload.model_validate( { "variant": ComposerVariant.WORKFLOW.value, @@ -1308,6 +2364,7 @@ def test_node_job_only_updates_inline_agent_soul(monkeypatch: pytest.MonkeyPatch tenant_id="tenant-1", agent_id="inline-agent-1", version=1, + home_snapshot_id="home-inline-1", config_snapshot='{"prompt":{"system_prompt":"old"}}', ) next_snapshot = AgentConfigSnapshot( @@ -1315,6 +2372,7 @@ def test_node_job_only_updates_inline_agent_soul(monkeypatch: pytest.MonkeyPatch tenant_id="tenant-1", agent_id="inline-agent-1", version=2, + home_snapshot_id="home-inline-2", config_snapshot=AgentSoulConfig.model_validate( { "model": { @@ -1334,6 +2392,7 @@ def test_node_job_only_updates_inline_agent_soul(monkeypatch: pytest.MonkeyPatch account_id=None, draft_owner_key="", base_snapshot_id="inline-version-1", + home_snapshot_id="home-inline-1", config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "old"}}), ) @@ -1387,6 +2446,7 @@ def test_node_job_only_updates_inline_agent_soul(monkeypatch: pytest.MonkeyPatch assert inline_agent.updated_by == "account-1" assert normal_draft.id == "draft-1" assert normal_draft.base_snapshot_id == "inline-version-2" + assert normal_draft.home_snapshot_id == "home-inline-2" assert normal_draft.config_snapshot_dict == next_snapshot.config_snapshot_dict assert normal_draft.updated_by == "account-1" @@ -1413,6 +2473,7 @@ def test_get_or_create_normal_agent_draft_rebases_stale_workflow_only_draft(): account_id=None, draft_owner_key="", base_snapshot_id="inline-version-1", + home_snapshot_id="home-inline-1", config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "old"}}), ) active_snapshot = AgentConfigSnapshot( @@ -1420,6 +2481,7 @@ def test_get_or_create_normal_agent_draft_rebases_stale_workflow_only_draft(): tenant_id="tenant-1", agent_id=agent.id, version=2, + home_snapshot_id="home-inline-2", config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "new"}}), ) session = FakeSession(scalar=[draft, active_snapshot]) @@ -1434,6 +2496,7 @@ def test_get_or_create_normal_agent_draft_rebases_stale_workflow_only_draft(): assert resolved is draft assert resolved.id == "draft-1" assert resolved.base_snapshot_id == "inline-version-2" + assert resolved.home_snapshot_id == "home-inline-2" assert resolved.config_snapshot_dict == active_snapshot.config_snapshot_dict assert resolved.updated_by == "account-2" assert session.flushes == 1 @@ -2098,11 +3161,14 @@ def test_drive_copy_scopes_include_declared_output_benchmark_files(): assert prefixes == {"tender-analyzer/"} -def test_composer_create_agents_syncs_active_config_has_model(monkeypatch: pytest.MonkeyPatch): +def test_composer_create_agents_syncs_active_config_has_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: fake_session = FakeSession() session = fake_session created_apps = [] hidden_backing_apps = [] + create_initial = MagicMock(return_value=SimpleNamespace(id="home-initial", snapshot_ref="backend-home-initial")) backing_agent = Agent( id="roster-agent-1", tenant_id="tenant-1", @@ -2133,11 +3199,17 @@ def test_composer_create_agents_syncs_active_config_has_model(monkeypatch: pytes monkeypatch.setattr(composer_service, "AppService", FakeAppService) monkeypatch.setattr(composer_service, "AgentRosterService", FakeAgentRosterService) + monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial) monkeypatch.setattr(AgentComposerService, "_require_account", lambda **kwargs: SimpleNamespace(id="account-1")) monkeypatch.setattr( AgentComposerService, "_require_version", - lambda **kwargs: SimpleNamespace(id="empty-version-1", tenant_id="tenant-1", agent_id="roster-agent-1"), + lambda **kwargs: SimpleNamespace( + id="empty-version-1", + tenant_id="tenant-1", + agent_id="roster-agent-1", + home_snapshot_id="home-roster-initial", + ), ) monkeypatch.setattr( AgentComposerService, @@ -2167,6 +3239,11 @@ def test_composer_create_agents_syncs_active_config_has_model(monkeypatch: pytes assert workflow_agent.active_config_snapshot_id == "version-with-model" assert workflow_agent.active_config_has_model is True assert workflow_agent.backing_app_id == "hidden-app-1" + create_initial.assert_called_once_with( + session=session, + tenant_id="tenant-1", + agent_id=workflow_agent.id, + ) assert hidden_backing_apps[0]["name"] == "Workflow Agent node-1" assert roster_agent.active_config_snapshot_id == "version-with-model" assert roster_agent.active_config_has_model is True @@ -2193,7 +3270,7 @@ def test_composer_require_account_raises_when_missing(monkeypatch: pytest.Monkey AgentComposerService._require_account(session=session, account_id="missing-account") -def test_composer_create_roster_agent_rolls_back_name_conflict(monkeypatch: pytest.MonkeyPatch): +def test_composer_create_roster_agent_maps_name_conflict_without_owning_rollback(monkeypatch: pytest.MonkeyPatch): fake_session = FakeSession() session = fake_session @@ -2215,7 +3292,7 @@ def test_composer_create_roster_agent_rolls_back_name_conflict(monkeypatch: pyte version_note=None, ) - assert fake_session.rollbacks == 1 + assert fake_session.rollbacks == 0 def test_composer_create_roster_agent_raises_when_backing_agent_missing(monkeypatch: pytest.MonkeyPatch): @@ -2257,7 +3334,7 @@ def test_agent_app_draft_match_does_not_mark_create_version_as_published(monkeyp source=AgentSource.AGENT_APP, active_config_snapshot_id="snapshot-1", ) - snapshot = SimpleNamespace(config_snapshot_dict=agent_soul) + snapshot = SimpleNamespace(config_snapshot_dict=agent_soul, home_snapshot_id="home-1") fake_session = FakeSession() session = fake_session monkeypatch.setattr(AgentComposerService, "_get_version_if_present", lambda **kwargs: snapshot) @@ -2268,6 +3345,7 @@ def test_agent_app_draft_match_does_not_mark_create_version_as_published(monkeyp tenant_id="tenant-1", agent=agent, agent_soul=agent_soul, + home_snapshot_id="home-1", ) is False ) @@ -2281,7 +3359,7 @@ def test_agent_app_draft_match_marks_publish_visible_revision_as_published(monke source=AgentSource.AGENT_APP, active_config_snapshot_id="snapshot-1", ) - snapshot = SimpleNamespace(config_snapshot_dict=agent_soul) + snapshot = SimpleNamespace(config_snapshot_dict=agent_soul, home_snapshot_id="home-1") fake_session = FakeSession(scalar=["publish-revision-1"]) session = fake_session monkeypatch.setattr(AgentComposerService, "_get_version_if_present", lambda **kwargs: snapshot) @@ -2292,6 +3370,7 @@ def test_agent_app_draft_match_marks_publish_visible_revision_as_published(monke tenant_id="tenant-1", agent=agent, agent_soul=agent_soul, + home_snapshot_id="home-1", ) is True ) @@ -2314,7 +3393,6 @@ def test_composer_version_helpers_and_lookup_errors(monkeypatch: pytest.MonkeyPa ) session = fake_session agent_soul = AgentSoulConfig.model_validate({"prompt": {"system_prompt": "new"}}) - version = AgentComposerService._create_config_version( session=session, tenant_id="tenant-1", @@ -2323,6 +3401,7 @@ def test_composer_version_helpers_and_lookup_errors(monkeypatch: pytest.MonkeyPa agent_soul=agent_soul, operation=AgentConfigRevisionOperation.SAVE_NEW_VERSION, version_note="note", + home_snapshot_id="home-1", ) updated_snapshot = AgentComposerService._update_current_version( session=session, @@ -2331,6 +3410,7 @@ def test_composer_version_helpers_and_lookup_errors(monkeypatch: pytest.MonkeyPa tenant_id="tenant-1", agent_id="agent-1", version=1, + home_snapshot_id="home-1", config_snapshot='{"prompt":{"system_prompt":"old"}}', ), account_id="account-1", @@ -2367,6 +3447,8 @@ def test_composer_version_helpers_and_lookup_errors(monkeypatch: pytest.MonkeyPa assert version.version == 2 assert updated_snapshot.version == 3 + assert version.home_snapshot_id == "home-1" + assert updated_snapshot.home_snapshot_id == "home-1" assert workflow.id == "workflow-1" @@ -2387,6 +3469,7 @@ def test_composer_current_version_and_error_paths(monkeypatch: pytest.MonkeyPatc tenant_id="tenant-1", agent_id="agent-1", version=1, + home_snapshot_id="home-1", config_snapshot='{"prompt":{"system_prompt":"old"}}', ) monkeypatch.setattr(AgentComposerService, "_require_version", lambda **kwargs: version) @@ -2395,7 +3478,6 @@ def test_composer_current_version_and_error_paths(monkeypatch: pytest.MonkeyPatc "_require_agent", lambda **kwargs: SimpleNamespace(updated_by=None, active_config_is_published=False), ) - result = AgentComposerService._save_to_current_version( session=session, tenant_id="tenant-1", @@ -2406,6 +3488,8 @@ def test_composer_current_version_and_error_paths(monkeypatch: pytest.MonkeyPatc assert result.updated_by == "account-1" assert result.current_snapshot_id != "version-1" + created_version = next(item for item in fake_session.added if isinstance(item, AgentConfigSnapshot)) + assert created_version.home_snapshot_id == "home-1" with pytest.raises(ValueError): AgentComposerService._require_binding(None) with pytest.raises(ValueError): @@ -2755,7 +3839,7 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat ) fake_session = FakeSession( scalar=["visible-revision"], - scalars=[[listed_version, older_listed_version], [older_listed_version, listed_version], [revision]], + scalars=[[], [listed_version, older_listed_version], [older_listed_version, listed_version], [revision]], ) agent = Agent( id="agent-1", @@ -2771,6 +3855,8 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat version.created_at = datetime(2026, 1, 4, 3, 4, 5, tzinfo=UTC) service = AgentRosterService(fake_session) + retire_snapshots = MagicMock(return_value=[]) + monkeypatch.setattr(AgentHomeSnapshotService, "retire_all_for_agent", retire_snapshots) monkeypatch.setattr(service, "_get_agent", lambda **kwargs: agent) monkeypatch.setattr(service, "_get_version", lambda **kwargs: version) monkeypatch.setattr( @@ -2791,6 +3877,7 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat assert updated["description"] == "new" assert agent.status == AgentStatus.ARCHIVED + retire_snapshots.assert_called_once_with(session=fake_session, tenant_id="tenant-1", agent_id="agent-1") assert versions[0]["id"] == "version-4" assert versions[0]["version"] == 2 assert versions[0]["display_version"] == 2 @@ -2807,6 +3894,63 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat assert detail["revisions"][0]["created_at"] == int(revision_created_at.timestamp()) +def test_roster_archive_retires_then_commits_before_enqueue(monkeypatch: pytest.MonkeyPatch) -> None: + session = FakeSession(scalars=[[SimpleNamespace(id="binding-1")]]) + service = AgentRosterService(session) + agent = SimpleNamespace( + status=AgentStatus.ACTIVE, + archived_by=None, + archived_at=None, + updated_by=None, + ) + events: list[str] = [] + monkeypatch.setattr(service, "_get_agent", lambda **_kwargs: agent) + monkeypatch.setattr( + AgentWorkspaceService, + "retire_binding", + MagicMock(side_effect=lambda **_kwargs: events.append("retire-binding") or "binding-1"), + ) + monkeypatch.setattr( + AgentHomeSnapshotService, + "retire_all_for_agent", + MagicMock(side_effect=lambda **_kwargs: events.append("retire-home") or ["home-1"]), + ) + session.commit = lambda: events.append("commit") # type: ignore[method-assign] + monkeypatch.setattr( + roster_service, + "enqueue_agent_resource_collection", + MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), + ) + + service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") + + assert events == ["retire-binding", "retire-home", "commit", "enqueue"] + + +def test_roster_archive_commit_failure_does_not_enqueue(monkeypatch: pytest.MonkeyPatch) -> None: + session = FakeSession(scalars=[[]]) + service = AgentRosterService(session) + monkeypatch.setattr( + service, + "_get_agent", + lambda **_kwargs: SimpleNamespace( + status=AgentStatus.ACTIVE, + archived_by=None, + archived_at=None, + updated_by=None, + ), + ) + monkeypatch.setattr(AgentHomeSnapshotService, "retire_all_for_agent", MagicMock(return_value=["home-1"])) + session.commit = MagicMock(side_effect=RuntimeError("commit failed")) # type: ignore[method-assign] + enqueue_collection = MagicMock() + monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(RuntimeError, match="commit failed"): + service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") + + enqueue_collection.assert_not_called() + + def test_roster_create_detail_and_lookup_helpers(monkeypatch: pytest.MonkeyPatch): fake_session = FakeSession( scalar=[ @@ -2818,6 +3962,13 @@ def test_roster_create_detail_and_lookup_helpers(monkeypatch: pytest.MonkeyPatch scalars=[[AgentConfigSnapshot(id="version-1", agent_id="agent-1", version=1)]], ) service = AgentRosterService(fake_session) + create_initial = MagicMock( + side_effect=[ + SimpleNamespace(id="home-roster", snapshot_ref="backend-home-roster"), + SimpleNamespace(id="home-backing", snapshot_ref="backend-home-backing"), + ] + ) + monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial) monkeypatch.setattr( AgentRosterService, "_get_or_create_agent_app_debug_conversation", @@ -2855,6 +4006,10 @@ def test_roster_create_detail_and_lookup_helpers(monkeypatch: pytest.MonkeyPatch assert created.role == "Research assistant" assert created.source == AgentSource.ROSTER assert created.active_config_snapshot_id is not None + assert create_initial.call_count == 2 + assert [ + snapshot.home_snapshot_id for snapshot in fake_session.added if isinstance(snapshot, AgentConfigSnapshot) + ] == ["home-roster", "home-backing"] assert created.active_config_has_model is False assert backing_agent.role == "Support agent" assert backing_agent.active_config_snapshot_id is not None @@ -2900,7 +4055,7 @@ def test_get_agent_runtime_app_model_creates_hidden_backing_app_for_existing_inl assert created_app.enable_api is False -def test_agent_app_debug_conversation_create_reuse_and_recreate(): +def test_agent_app_build_conversation_create_reuse_and_recreate(): agent = Agent( id="agent-1", tenant_id="tenant-1", @@ -2914,7 +4069,7 @@ def test_agent_app_debug_conversation_create_reuse_and_recreate(): ) create_session = FakeSession(scalar=[agent, None]) - created_id = AgentRosterService(create_session).get_or_create_agent_app_debug_conversation_id( + created_id = AgentRosterService(create_session).get_or_create_build_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", @@ -2939,7 +4094,7 @@ def test_agent_app_debug_conversation_create_reuse_and_recreate(): conversation_id="existing-conversation", ) reuse_session = FakeSession(scalar=[agent, existing_mapping, "existing-conversation"]) - reused_id = AgentRosterService(reuse_session).get_or_create_agent_app_debug_conversation_id( + reused_id = AgentRosterService(reuse_session).get_or_create_build_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", @@ -2957,7 +4112,7 @@ def test_agent_app_debug_conversation_create_reuse_and_recreate(): conversation_id="deleted-conversation", ) recreate_session = FakeSession(scalar=[agent, stale_mapping, None]) - recreated_id = AgentRosterService(recreate_session).get_or_create_agent_app_debug_conversation_id( + recreated_id = AgentRosterService(recreate_session).get_or_create_build_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", @@ -2983,17 +4138,15 @@ def test_agent_app_debug_conversations_are_isolated_by_draft_type(): session = FakeSession(scalar=[agent, None, agent, None]) service = AgentRosterService(session) - build_conversation_id = service.get_or_create_agent_app_debug_conversation_id( + build_conversation_id = service.get_or_create_build_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", - draft_type=AgentConfigDraftType.DEBUG_BUILD, ) - preview_conversation_id = service.get_or_create_agent_app_debug_conversation_id( + preview_conversation_id = service.rotate_preview_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", - draft_type=AgentConfigDraftType.DRAFT, ) mappings = [value for value in session.added if isinstance(value, AgentDebugConversation)] @@ -3034,7 +4187,7 @@ def test_agent_app_debug_conversation_requires_app_binding(): ) -def test_load_or_create_agent_app_debug_conversations_supports_runtime_backed_agents(): +def test_load_or_create_build_conversations_supports_runtime_backed_agents(): valid_agent = Agent( id="agent-1", tenant_id="tenant-1", @@ -3067,7 +4220,7 @@ def test_load_or_create_agent_app_debug_conversations_supports_runtime_backed_ag ) fake_session = FakeSession(scalar=[None]) - result = AgentRosterService(fake_session).load_or_create_agent_app_debug_conversation_ids_by_agent_id( + result = AgentRosterService(fake_session).load_or_create_build_conversation_ids_by_agent_id( tenant_id="tenant-1", agents=[valid_agent, wrong_tenant_agent, workflow_agent], account_id="account-1", @@ -3121,6 +4274,7 @@ def test_restore_roster_agent_version_switches_active_snapshot(monkeypatch: pyte tenant_id="tenant-1", agent_id="agent-1", version=2, + home_snapshot_id="home-version-2", config_snapshot=_agent_soul_with_model(), ) @@ -3149,6 +4303,7 @@ def test_restore_roster_agent_version_switches_active_snapshot(monkeypatch: pyte assert draft.agent_id == "agent-1" assert draft.draft_type == AgentConfigDraftType.DRAFT assert draft.base_snapshot_id == "version-2" + assert draft.home_snapshot_id == "home-version-2" assert draft.config_snapshot_dict == _agent_soul_with_model().model_dump(mode="json") assert draft.updated_by == "account-1" @@ -3395,9 +4550,11 @@ class TestAgentAppBackingAgent: ``Agent.app_id``. ``AppService.create_app`` builds the backing agent inside its own transaction, so the helper must add+flush without committing.""" - def test_create_backing_agent_for_app_links_app_and_seeds_default_soul(self): + def test_create_backing_agent_for_app_links_app_and_seeds_default_soul(self, monkeypatch: pytest.MonkeyPatch): session = FakeSession() service = AgentRosterService(session) + create_initial = MagicMock(return_value=SimpleNamespace(id="home-1", snapshot_ref="backend-home-1")) + monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial) agent = service.create_backing_agent_for_app( tenant_id="tenant-1", @@ -3420,6 +4577,12 @@ class TestAgentAppBackingAgent: snapshots = [a for a in session.added if isinstance(a, AgentConfigSnapshot)] assert len(snapshots) == 1 assert snapshots[0].version == 1 + assert snapshots[0].home_snapshot_id == "home-1" + create_initial.assert_called_once_with( + session=session, + tenant_id="tenant-1", + agent_id=agent.id, + ) assert agent.active_config_snapshot_id == snapshots[0].id revisions = [ a for a in session.added if getattr(a, "operation", None) == AgentConfigRevisionOperation.CREATE_VERSION @@ -3482,7 +4645,7 @@ class TestAgentAppBackingAgent: with pytest.raises(roster_service.AgentNotFoundError): service.get_agent_app_model(tenant_id="tenant-1", agent_id="agent-x") - def test_refresh_agent_app_debug_conversation_creates_mapping(self): + def test_reset_build_conversation_creates_mapping(self): agent = Agent( id="agent-1", tenant_id="tenant-1", @@ -3494,10 +4657,10 @@ class TestAgentAppBackingAgent: status=AgentStatus.ACTIVE, app_id="app-1", ) - session = FakeSession(scalar=[agent, None]) + session = FakeSession(scalar=[agent, None, None]) service = AgentRosterService(session) - conversation_id = service.refresh_agent_app_debug_conversation_id( + conversation_id = service.reset_build_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", @@ -3520,7 +4683,7 @@ class TestAgentAppBackingAgent: assert session.deleted == [] assert session.commits == 1 - def test_refresh_agent_app_debug_conversation_replaces_existing_mapping(self): + def test_rotate_preview_conversation_retires_exact_binding(self, monkeypatch: pytest.MonkeyPatch): agent = Agent( id="agent-1", tenant_id="tenant-1", @@ -3533,10 +4696,32 @@ class TestAgentAppBackingAgent: app_id="app-1", ) mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") - session = FakeSession(scalar=[agent, mapping]) + previous_conversation = SimpleNamespace( + id="old-conversation", + agent_workspace_binding_id="binding-1", + ) + binding = SimpleNamespace(agent_id=agent.id) + session = FakeSession(scalar=[agent, mapping, previous_conversation]) service = AgentRosterService(session) + events: list[str] = [] + original_commit = session.commit - conversation_id = service.refresh_agent_app_debug_conversation_id( + def commit() -> None: + original_commit() + events.append("commit") + + session.commit = commit # type: ignore[method-assign] + get_active_binding = MagicMock(return_value=binding) + retire_binding = MagicMock(side_effect=lambda **_kwargs: events.append("retire") or "binding-1") + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + monkeypatch.setattr( + roster_service, + "enqueue_agent_resource_collection", + MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), + ) + + conversation_id = service.rotate_preview_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", @@ -3550,10 +4735,17 @@ class TestAgentAppBackingAgent: assert conversations[0].id == conversation_id assert session.deleted == [] assert session.commits == 1 + assert events == ["retire", "commit", "enqueue"] + owner_scope = get_active_binding.call_args.kwargs["expected_owner_scope"] + assert owner_scope.owner_type is AgentWorkspaceOwnerType.CONVERSATION + assert owner_scope.owner_id == previous_conversation.id + retire_binding.assert_called_once_with( + session=session, + tenant_id="tenant-1", + binding_id="binding-1", + ) - def test_refresh_agent_app_debug_conversation_rotates_preview_mapping_each_time( - self, monkeypatch: pytest.MonkeyPatch - ): + def test_reset_build_conversation_retires_build_draft_binding(self, monkeypatch: pytest.MonkeyPatch): agent = Agent( id="agent-1", tenant_id="tenant-1", @@ -3565,125 +4757,38 @@ class TestAgentAppBackingAgent: status=AgentStatus.ACTIVE, app_id="app-1", ) - session = FakeSession(scalar=[agent, None]) - service = AgentRosterService(session) - cleanup = MagicMock() - monkeypatch.setattr(service, "_cleanup_debug_conversation_runtime_sessions", cleanup) + mapping = SimpleNamespace(app_id="app-1", conversation_id="old-build-conversation") + build_draft = SimpleNamespace(id="build-draft-1", agent_workspace_binding_id="binding-1") + session = FakeSession(scalar=[agent, mapping, build_draft]) + get_active_binding = MagicMock(return_value=SimpleNamespace(agent_id=agent.id)) + retire_binding = MagicMock(return_value="binding-1") + enqueue_collection = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", enqueue_collection) - first_conversation_id = service.refresh_agent_app_debug_conversation_id( + conversation_id = AgentRosterService(session).reset_build_conversation( tenant_id="tenant-1", - agent_id="agent-1", + agent_id=agent.id, account_id="account-1", - draft_type=AgentConfigDraftType.DRAFT, - ) - mapping = next(value for value in session.added if isinstance(value, AgentDebugConversation)) - session._scalar.extend([agent, mapping]) - second_conversation_id = service.refresh_agent_app_debug_conversation_id( - tenant_id="tenant-1", - agent_id="agent-1", - account_id="account-1", - draft_type=AgentConfigDraftType.DRAFT, ) - assert first_conversation_id != second_conversation_id - assert mapping.draft_type == AgentConfigDraftType.DRAFT - assert mapping.conversation_id == second_conversation_id - assert len([value for value in session.added if isinstance(value, Conversation)]) == 2 - cleanup.assert_called_once_with( - tenant_id="tenant-1", - agent_id="agent-1", - account_id="account-1", - draft_type=AgentConfigDraftType.DRAFT, - app_id="app-1", - conversation_id=first_conversation_id, - ) - - @pytest.mark.parametrize( - "draft_type", - [AgentConfigDraftType.DRAFT, AgentConfigDraftType.DEBUG_BUILD], - ) - def test_refresh_agent_app_debug_conversation_enqueues_cleanup_for_old_runtime_sessions( - self, - monkeypatch: pytest.MonkeyPatch, - draft_type: AgentConfigDraftType, - ): - agent = Agent( - id="agent-1", - tenant_id="tenant-1", - name="Iris", - description="", - agent_kind=AgentKind.DIFY_AGENT, - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - status=AgentStatus.ACTIVE, - app_id="app-1", - ) - mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") - stored_session = SimpleNamespace( - scope=SimpleNamespace( - tenant_id="tenant-1", - app_id="old-app", - conversation_id="old-conversation", - agent_id="agent-9", - agent_config_snapshot_id="snap-9", - ), - session_snapshot=CompositorSessionSnapshot(layers=[]), - backend_run_id="run-old", - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - ) - session = FakeSession(scalar=[agent, mapping]) - service = AgentRosterService(session) - events: list[str] = [] - original_commit = session.commit - - def record_commit() -> None: - events.append("commit") - original_commit() - - def list_active_sessions(**kwargs: object) -> list[object]: - events.append("cleanup") - return [stored_session] - - cleanup_delay = MagicMock() - cleanup_store = MagicMock() - cleanup_store.list_active_sessions_for_conversation.side_effect = list_active_sessions - monkeypatch.setattr(session, "commit", record_commit) - monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store) - monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) - - conversation_id = service.refresh_agent_app_debug_conversation_id( - tenant_id="tenant-1", - agent_id="agent-1", - account_id="account-1", - draft_type=draft_type, - ) - - cleanup_store.list_active_sessions_for_conversation.assert_called_once_with( - tenant_id="tenant-1", - app_id="old-app", - conversation_id="old-conversation", - ) - cleanup_delay.assert_called_once() - payload = cleanup_delay.call_args.args[0] - assert payload["metadata"]["conversation_id"] == "old-conversation" - assert payload["metadata"]["agent_id"] == "agent-9" - assert payload["metadata"]["draft_type"] == draft_type.value - assert ( - payload["idempotency_key"] - == f"tenant-1:agent-1:account-1:{draft_type.value}:old-conversation:debug-session-cleanup:" - "agent-9:snap-9:run-old" - ) - cleanup_store.mark_cleaned.assert_called_once_with( - scope=stored_session.scope, - backend_run_id="run-old", - ) - assert mapping.app_id == "app-1" assert mapping.conversation_id == conversation_id - assert events == ["commit", "cleanup"] + assert build_draft.agent_workspace_binding_id is None + owner_scope = get_active_binding.call_args.kwargs["expected_owner_scope"] + assert owner_scope.owner_type is AgentWorkspaceOwnerType.BUILD_DRAFT + assert owner_scope.owner_id == build_draft.id + retire_binding.assert_called_once_with( + session=session, + tenant_id="tenant-1", + binding_id="binding-1", + ) + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=("binding-1",), + ) - def test_refresh_agent_app_debug_conversation_does_not_cleanup_when_commit_fails( - self, monkeypatch: pytest.MonkeyPatch - ): + def test_preview_rotation_commit_failure_rolls_back_before_enqueue(self, monkeypatch: pytest.MonkeyPatch): agent = Agent( id="agent-1", tenant_id="tenant-1", @@ -3695,29 +4800,38 @@ class TestAgentAppBackingAgent: status=AgentStatus.ACTIVE, app_id="app-1", ) - mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") - session = FakeSession(scalar=[agent, mapping]) - service = AgentRosterService(session) - runtime_store_factory = MagicMock() - cleanup_delay = MagicMock() - monkeypatch.setattr(session, "commit", MagicMock(side_effect=RuntimeError("database unavailable"))) - monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", runtime_store_factory) - monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) + mapping = SimpleNamespace(app_id="app-1", conversation_id="old-conversation") + previous_conversation = SimpleNamespace( + id="old-conversation", + agent_workspace_binding_id="binding-1", + ) + session = FakeSession(scalar=[agent, mapping, previous_conversation]) - with pytest.raises(RuntimeError, match="database unavailable"): - service.refresh_agent_app_debug_conversation_id( + def fail_commit() -> None: + session.commits += 1 + raise RuntimeError("commit failed") + + session.commit = fail_commit # type: ignore[method-assign] + monkeypatch.setattr( + AgentWorkspaceService, + "get_active_binding", + MagicMock(return_value=SimpleNamespace(agent_id=agent.id)), + ) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", MagicMock(return_value="binding-1")) + enqueue_collection = MagicMock() + monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(RuntimeError, match="commit failed"): + AgentRosterService(session).rotate_preview_conversation( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", - draft_type=AgentConfigDraftType.DRAFT, ) - runtime_store_factory.assert_not_called() - cleanup_delay.assert_not_called() + assert session.rollbacks == 1 + enqueue_collection.assert_not_called() - def test_refresh_agent_app_debug_conversation_marks_old_runtime_sessions_clean_when_enqueue_fails( - self, monkeypatch: pytest.MonkeyPatch - ): + def test_build_reset_commit_failure_rolls_back_before_enqueue(self, monkeypatch: pytest.MonkeyPatch): agent = Agent( id="agent-1", tenant_id="tenant-1", @@ -3729,83 +4843,48 @@ class TestAgentAppBackingAgent: status=AgentStatus.ACTIVE, app_id="app-1", ) - mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") - stored_session = SimpleNamespace( - scope=SimpleNamespace( + mapping = SimpleNamespace(app_id="app-1", conversation_id="old-build-conversation") + build_draft = SimpleNamespace(id="build-draft-1", agent_workspace_binding_id="binding-1") + session = FakeSession(scalar=[agent, mapping, build_draft]) + events: list[str] = [] + original_rollback = session.rollback + + def fail_commit() -> None: + assert build_draft.agent_workspace_binding_id is None + events.append("commit") + raise RuntimeError("commit failed") + + def rollback() -> None: + events.append("rollback") + original_rollback() + + session.commit = fail_commit # type: ignore[method-assign] + session.rollback = rollback # type: ignore[method-assign] + get_active_binding = MagicMock(return_value=SimpleNamespace(agent_id=agent.id)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) + retire_binding = MagicMock(side_effect=lambda **_kwargs: events.append("retire") or "binding-1") + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + enqueue_collection = MagicMock() + monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(RuntimeError, match="commit failed"): + AgentRosterService(session).reset_build_conversation( tenant_id="tenant-1", - app_id="old-app", - conversation_id="old-conversation", - agent_id="agent-9", - agent_config_snapshot_id="snap-9", - ), - session_snapshot=CompositorSessionSnapshot(layers=[]), - backend_run_id="run-old", - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - ) - session = FakeSession(scalar=[agent, mapping]) - service = AgentRosterService(session) - cleanup_store = MagicMock() - cleanup_store.list_active_sessions_for_conversation.return_value = [stored_session] - monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store) - monkeypatch.setattr( - roster_service.cleanup_conversation_agent_runtime_session, - "delay", - MagicMock(side_effect=RuntimeError("queue down")), - ) + agent_id=agent.id, + account_id="account-1", + ) - service.refresh_agent_app_debug_conversation_id( + assert events == ["retire", "commit", "rollback"] + assert session.rollbacks == 1 + owner_scope = get_active_binding.call_args.kwargs["expected_owner_scope"] + assert owner_scope.owner_type is AgentWorkspaceOwnerType.BUILD_DRAFT + assert owner_scope.owner_id == build_draft.id + retire_binding.assert_called_once_with( + session=session, tenant_id="tenant-1", - agent_id="agent-1", - account_id="account-1", + binding_id="binding-1", ) - - cleanup_store.mark_cleaned.assert_called_once_with( - scope=stored_session.scope, - backend_run_id="run-old", - ) - - def test_refresh_agent_app_debug_conversation_ignores_mark_cleaned_failure(self, monkeypatch: pytest.MonkeyPatch): - agent = Agent( - id="agent-1", - tenant_id="tenant-1", - name="Iris", - description="", - agent_kind=AgentKind.DIFY_AGENT, - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - status=AgentStatus.ACTIVE, - app_id="app-1", - ) - mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") - stored_session = SimpleNamespace( - scope=SimpleNamespace( - tenant_id="tenant-1", - app_id="old-app", - conversation_id="old-conversation", - agent_id="agent-9", - agent_config_snapshot_id="snap-9", - ), - session_snapshot=CompositorSessionSnapshot(layers=[]), - backend_run_id="run-old", - runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], - ) - session = FakeSession(scalar=[agent, mapping]) - service = AgentRosterService(session) - cleanup_store = MagicMock() - cleanup_store.list_active_sessions_for_conversation.return_value = [stored_session] - cleanup_store.mark_cleaned.side_effect = RuntimeError("cleanup bookkeeping failed") - monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store) - monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", MagicMock()) - - conversation_id = service.refresh_agent_app_debug_conversation_id( - tenant_id="tenant-1", - agent_id="agent-1", - account_id="account-1", - ) - - assert mapping.app_id == "app-1" - assert mapping.conversation_id == conversation_id - assert session.commits == 1 + enqueue_collection.assert_not_called() def test_duplicate_agent_app_copies_app_config_and_active_soul(self, monkeypatch: pytest.MonkeyPatch): source_config = SimpleNamespace( @@ -3871,51 +4950,54 @@ class TestAgentAppBackingAgent: active_config_snapshot_id="source-version", active_config_has_model=True, ) - target_agent = Agent( - id="target-agent", - tenant_id="tenant-1", - name="Iris copy", - description="source desc", - role="Analyst", - agent_kind=AgentKind.DIFY_AGENT, - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - status=AgentStatus.ACTIVE, - app_id="target-app", - active_config_snapshot_id="target-version", - ) source_version = AgentConfigSnapshot( id="source-version", tenant_id="tenant-1", agent_id="source-agent", version=1, + home_snapshot_id="home-source", config_snapshot=_agent_soul_with_model(), summary="configured", version_note="v1", created_by="account-1", ) - target_version = AgentConfigSnapshot( - id="target-version", - tenant_id="tenant-1", - agent_id="target-agent", - version=1, - config_snapshot=AgentSoulConfig(), - created_by="account-1", - ) session = FakeSession( - scalar=[source_agent, source_app, source_agent, target_agent, source_version, target_version], + scalar=[source_agent, source_app, source_agent], scalars=[[]], ) captured: dict[str, object] = {} + create_initial = MagicMock(return_value=SimpleNamespace(id="home-target", snapshot_ref="backend-home-target")) class FakeAppService: def create_app(self, tenant_id: str, params, account: object, *, session) -> object: captured["tenant_id"] = tenant_id captured["params"] = params captured["account"] = account + target_agent = AgentRosterService(session).create_backing_agent_for_app( + tenant_id=tenant_id, + account_id="account-1", + app_id=target_app.id, + name=params.name, + description=params.description, + role=params.agent_role, + ) + target_version = next( + value + for value in session.added + if isinstance(value, AgentConfigSnapshot) and value.agent_id == target_agent.id + ) + captured["target_agent"] = target_agent + captured["target_version"] = target_version + session._scalar.extend([target_agent, source_version, target_version]) return target_app monkeypatch.setattr(roster_service, "AppService", FakeAppService) + monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial) + monkeypatch.setattr( + AgentRosterService, + "_get_or_create_agent_app_debug_conversation", + lambda _self, **_kwargs: None, + ) monkeypatch.setattr( roster_service.FeatureService, "get_system_features", @@ -3941,11 +5023,20 @@ class TestAgentAppBackingAgent: assert target_config.opening_statement == "hello" assert target_config.file_upload == '{"image": {"enabled": true}}' assert target_config.updated_by == "account-1" + target_agent = captured["target_agent"] + target_version = captured["target_version"] assert target_version.config_snapshot.model.model == "gpt-4o" + assert source_version.home_snapshot_id == "home-source" + assert target_version.home_snapshot_id == "home-target" assert target_version.summary == "configured" assert target_version.version_note == "v1" assert target_agent.active_config_has_model is True assert target_agent.updated_by == "account-1" + create_initial.assert_called_once_with( + session=session, + tenant_id="tenant-1", + agent_id=target_agent.id, + ) assert session.commits == 1 def test_duplicate_agent_app_inherits_webapp_access_mode(self, monkeypatch: pytest.MonkeyPatch): @@ -5015,35 +6106,210 @@ class TestWorkflowAgentDraftBindingSync: assert node_job.declared_outputs == [] assert existing_binding.current_snapshot_id == "snapshot-2" - def test_deletes_draft_binding_when_agent_node_removed(self): + @pytest.mark.parametrize( + "sqlite_session", + [(Agent, AgentConfigRevision, AgentConfigSnapshot, WorkflowAgentNodeBinding)], + indirect=True, + ) + def test_deletes_draft_binding_and_returns_only_replaced_inline_agent( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ) -> None: workflow = Workflow( id="workflow-1", tenant_id="tenant-1", app_id="app-1", version=Workflow.VERSION_DRAFT, - graph='{"nodes":[]}', + graph=json.dumps( + { + "nodes": [ + { + "id": "kept-node", + "data": { + "type": "agent", + "version": "2", + "agent_binding": { + "binding_type": "inline_agent", + "agent_id": "inline-kept", + "current_snapshot_id": "snapshot-kept", + }, + }, + }, + { + "id": "roster-node", + "data": { + "type": "agent", + "version": "2", + "agent_binding": { + "binding_type": "roster_agent", + "agent_id": "roster-new", + }, + }, + }, + { + "id": "inline-replaced-node", + "data": { + "type": "agent", + "version": "2", + "agent_binding": { + "binding_type": "inline_agent", + "agent_id": "inline-new", + "current_snapshot_id": "snapshot-inline-new", + }, + }, + }, + ] + } + ), ) - stale_binding = WorkflowAgentNodeBinding( - id="binding-1", + removed_inline = WorkflowAgentNodeBinding( + id="binding-inline-removed", tenant_id="tenant-1", app_id="app-1", workflow_id="workflow-1", workflow_version=Workflow.VERSION_DRAFT, node_id="removed-node", - binding_type=WorkflowAgentBindingType.ROSTER_AGENT, - agent_id="agent-1", - current_snapshot_id="snapshot-1", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="inline-removed", + current_snapshot_id="snapshot-removed", node_job_config=WorkflowNodeJobConfig(), ) - session = FakeSession(scalars=[[stale_binding]]) - - WorkflowAgentPublishService.sync_roster_agent_bindings_for_draft( - session=session, + kept_inline = WorkflowAgentNodeBinding( + id="binding-inline-kept", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version=Workflow.VERSION_DRAFT, + node_id="kept-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="inline-kept", + current_snapshot_id="snapshot-kept", + node_job_config=WorkflowNodeJobConfig(), + ) + removed_roster = WorkflowAgentNodeBinding( + id="binding-roster-removed", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version=Workflow.VERSION_DRAFT, + node_id="removed-roster-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="roster-removed", + current_snapshot_id="snapshot-roster", + node_job_config=WorkflowNodeJobConfig(), + ) + old_inline_replaced_by_roster = WorkflowAgentNodeBinding( + id="binding-inline-to-roster", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version=Workflow.VERSION_DRAFT, + node_id="roster-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="inline-old-roster", + current_snapshot_id="snapshot-inline-old-roster", + node_job_config=WorkflowNodeJobConfig(), + ) + old_inline_replaced_by_inline = WorkflowAgentNodeBinding( + id="binding-inline-to-inline", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version=Workflow.VERSION_DRAFT, + node_id="inline-replaced-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="inline-old-inline", + current_snapshot_id="snapshot-inline-old-inline", + node_job_config=WorkflowNodeJobConfig(), + ) + kept_agent = Agent( + id="inline-kept", + tenant_id="tenant-1", + name="Kept inline", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + status=AgentStatus.ACTIVE, + app_id="app-1", + workflow_id="workflow-1", + workflow_node_id="kept-node", + active_config_snapshot_id="snapshot-kept", + ) + replacement_agent = Agent( + id="inline-new", + tenant_id="tenant-1", + name="Replacement inline", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + status=AgentStatus.ACTIVE, + app_id="app-1", + workflow_id="workflow-1", + workflow_node_id="inline-replaced-node", + active_config_snapshot_id="snapshot-inline-new", + ) + roster_agent = Agent( + id="roster-new", + tenant_id="tenant-1", + name="Roster replacement", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.ROSTER, + status=AgentStatus.ACTIVE, + active_config_snapshot_id="snapshot-roster-new", + ) + kept_snapshot = AgentConfigSnapshot( + id="snapshot-kept", + tenant_id="tenant-1", + agent_id="inline-kept", + version=1, + config_snapshot=AgentSoulConfig(), + home_snapshot_id="home-kept", + ) + replacement_snapshot = AgentConfigSnapshot( + id="snapshot-inline-new", + tenant_id="tenant-1", + agent_id="inline-new", + version=1, + config_snapshot=AgentSoulConfig(), + home_snapshot_id="home-inline-new", + ) + sqlite_session.add_all( + [ + kept_agent, + replacement_agent, + roster_agent, + kept_snapshot, + replacement_snapshot, + removed_inline, + kept_inline, + removed_roster, + old_inline_replaced_by_roster, + old_inline_replaced_by_inline, + ] + ) + sqlite_session.commit() + retirement_candidates = WorkflowAgentPublishService.sync_agent_bindings_for_draft( + session=sqlite_session, draft_workflow=workflow, account_id="account-1", ) - assert session.deleted == [stale_binding] + assert sqlite_session.get(WorkflowAgentNodeBinding, removed_inline.id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, removed_roster.id) is None + assert kept_inline.binding_type == WorkflowAgentBindingType.INLINE_AGENT + assert kept_inline.agent_id == "inline-kept" + assert old_inline_replaced_by_roster.binding_type == WorkflowAgentBindingType.ROSTER_AGENT + assert old_inline_replaced_by_roster.agent_id == "roster-new" + assert old_inline_replaced_by_roster.current_snapshot_id == "snapshot-roster-new" + assert old_inline_replaced_by_inline.binding_type == WorkflowAgentBindingType.INLINE_AGENT + assert old_inline_replaced_by_inline.agent_id == "inline-new" + assert old_inline_replaced_by_inline.current_snapshot_id == "snapshot-inline-new" + assert retirement_candidates == {"inline-removed", "inline-old-roster", "inline-old-inline"} def test_dataset_rows_filters_malformed_ids(monkeypatch: pytest.MonkeyPatch): @@ -5156,7 +6422,9 @@ def test_save_agent_composer_allows_incomplete_knowledge_draft(monkeypatch: pyte active_config_is_published=True, updated_by=None, ) - active_version = SimpleNamespace(config_snapshot_dict=AgentSoulConfig().model_dump(mode="json")) + active_version = SimpleNamespace( + home_snapshot_id="home-initial", config_snapshot_dict=AgentSoulConfig().model_dump(mode="json") + ) fake_session = FakeSession(scalar=[agent]) saved = {} @@ -5170,7 +6438,7 @@ def test_save_agent_composer_allows_incomplete_knowledge_draft(monkeypatch: pyte monkeypatch.setattr( AgentComposerService, "_save_agent_draft", - lambda **kwargs: saved.update(kwargs) or SimpleNamespace(id="draft-1"), + lambda **kwargs: saved.update(kwargs) or SimpleNamespace(id="draft-1", home_snapshot_id="home-initial"), ) monkeypatch.setattr(AgentComposerService, "_get_version_if_present", lambda **_kwargs: active_version) monkeypatch.setattr(AgentComposerService, "load_agent_composer", lambda **_kwargs: {"loaded": True}) diff --git a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py new file mode 100644 index 00000000000..bf17dbfb169 --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py @@ -0,0 +1,181 @@ +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from sqlalchemy.orm import Session + +from models.agent import ( + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigSnapshot, + AgentHomeSnapshot, + AgentWorkingResourceStatus, +) +from models.agent_config_entities import AgentSoulConfig +from services.agent.errors import AgentBuildSandboxNotFoundError +from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.workspace_service import AgentWorkspaceService + + +def _build_draft() -> AgentConfigDraft: + return AgentConfigDraft( + id="build-1", + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + home_snapshot_id="home-old", + agent_workspace_binding_id="binding-1", + config_snapshot=AgentSoulConfig(), + ) + + +def _client(*, snapshot_ref: str = "snapshot-ref-1") -> MagicMock: + client = MagicMock() + client.initialize_home_snapshot_sync.return_value = SimpleNamespace(snapshot_ref=snapshot_ref) + client.create_home_snapshot_from_binding_sync.return_value = SimpleNamespace(snapshot_ref=snapshot_ref) + return client + + +def test_create_initial_persists_backend_snapshot_ref(monkeypatch: pytest.MonkeyPatch) -> None: + session = MagicMock() + client = _client() + monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client)) + + snapshot = AgentHomeSnapshotService.create_initial( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + ) + + assert snapshot.snapshot_ref == "snapshot-ref-1" + assert snapshot.status is AgentWorkingResourceStatus.ACTIVE + session.add.assert_called_once_with(snapshot) + session.flush.assert_called_once_with() + + +def test_create_initial_flush_failure_does_not_delete_backend_snapshot(monkeypatch: pytest.MonkeyPatch) -> None: + session = MagicMock() + session.flush.side_effect = RuntimeError("flush failed") + client = _client() + monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client)) + + with pytest.raises(RuntimeError, match="flush failed"): + AgentHomeSnapshotService.create_initial( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + ) + + client.delete_home_snapshot_sync.assert_not_called() + + +@pytest.mark.parametrize( + ("app_id", "backing_app_id", "expected_runtime_app_id"), + [ + ("app-1", None, "app-1"), + ("workflow-app-1", "runtime-app-1", "runtime-app-1"), + ], +) +def test_build_apply_checkpoints_exact_active_binding( + monkeypatch: pytest.MonkeyPatch, + app_id: str, + backing_app_id: str | None, + expected_runtime_app_id: str, +) -> None: + session = MagicMock() + session.scalar.return_value = SimpleNamespace(app_id=app_id, backing_app_id=backing_app_id) + binding = SimpleNamespace( + backend_binding_ref="binding-ref-1", + agent_id="agent-1", + base_home_snapshot_id="home-old", + agent_config_version_id="build-1", + agent_config_version_kind="build_draft", + ) + get_binding = MagicMock(return_value=binding) + client = _client(snapshot_ref="snapshot-ref-2") + monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding) + monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", MagicMock()) + + snapshot = AgentHomeSnapshotService.create_for_build_apply( + session=session, + build_draft=_build_draft(), + ) + + assert get_binding.call_args.kwargs["binding_id"] == "binding-1" + assert get_binding.call_args.kwargs["expected_owner_scope"].app_id == expected_runtime_app_id + request = client.create_home_snapshot_from_binding_sync.call_args.args[0] + assert request.backend_binding_ref == "binding-ref-1" + assert snapshot.snapshot_ref == "snapshot-ref-2" + + +def test_build_apply_fails_fast_without_source_binding() -> None: + session = MagicMock() + build_draft = _build_draft() + build_draft.agent_workspace_binding_id = None + + with pytest.raises(AgentBuildSandboxNotFoundError): + AgentHomeSnapshotService.create_for_build_apply( + session=session, + build_draft=build_draft, + ) + + +def test_home_snapshot_collection_database_failure_is_best_effort(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + session.scalar.side_effect = RuntimeError("database unavailable") + log_exception = MagicMock() + monkeypatch.setattr( + "services.agent.home_snapshot_service.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr("services.agent.home_snapshot_service.logger.exception", log_exception) + + AgentHomeSnapshotService.collect_retired_home_snapshot( + tenant_id="tenant-1", + home_snapshot_id="home-1", + ) + + session.scalar.assert_called_once() + log_exception.assert_called_once_with( + "Failed to collect retired Agent Home Snapshot", + extra={"tenant_id": "tenant-1", "home_snapshot_id": "home-1"}, + ) + + +@pytest.mark.parametrize( + "sqlite_session", + [(AgentHomeSnapshot, AgentConfigDraft, AgentConfigSnapshot)], + indirect=True, +) +def test_home_snapshot_collection_final_delete_failure_is_best_effort( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + snapshot = AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="snapshot-ref-1", + status=AgentWorkingResourceStatus.RETIRED, + ) + sqlite_session.add(snapshot) + sqlite_session.commit() + commit = MagicMock(side_effect=RuntimeError("database unavailable")) + delete = MagicMock() + monkeypatch.setattr( + "services.agent.home_snapshot_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + monkeypatch.setattr(sqlite_session, "commit", commit) + monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete) + + AgentHomeSnapshotService.collect_retired_home_snapshot( + tenant_id="tenant-1", + home_snapshot_id="home-1", + ) + + delete.assert_called_once_with(snapshot_ref="snapshot-ref-1") diff --git a/api/tests/unit_tests/services/agent/test_retirement_service.py b/api/tests/unit_tests/services/agent/test_retirement_service.py new file mode 100644 index 00000000000..71362eecfe3 --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_retirement_service.py @@ -0,0 +1,205 @@ +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from sqlalchemy.orm import Session + +from models.agent import ( + Agent, + AgentConfigVersionKind, + AgentHomeSnapshot, + AgentKind, + AgentScope, + AgentSource, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) +from models.enums import AppStatus +from models.model import App, AppMode +from models.workflow import Workflow, WorkflowType +from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.retirement_service import WorkflowAgentRetirementService +from services.agent.workspace_service import AgentWorkspaceService + + +def test_retire_unowned_commits_resource_retirement(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + session.scalars.return_value.all.return_value = [SimpleNamespace(id="binding-1")] + monkeypatch.setattr( + "services.agent.retirement_service.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr( + WorkflowAgentRetirementService, + "archive_unowned", + MagicMock(return_value=["agent-1"]), + ) + monkeypatch.setattr( + AgentWorkspaceService, + "retire_binding", + MagicMock(return_value="binding-1"), + ) + monkeypatch.setattr( + AgentHomeSnapshotService, + "retire_all_for_agent", + MagicMock(return_value=["home-1"]), + ) + + result = WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=["agent-1"], + account_id="account-1", + ) + + assert result == (["binding-1"], ["home-1"]) + session.commit.assert_called_once_with() + + +def _workflow_only_agent() -> Agent: + return Agent( + id="agent-1", + tenant_id="tenant-1", + name="Inline Agent", + description="", + role="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + status=AgentStatus.ACTIVE, + ) + + +@pytest.mark.parametrize( + "sqlite_session", + [(Agent, App, Workflow, WorkflowAgentNodeBinding)], + indirect=True, +) +def test_retire_unowned_keeps_effectively_owned_agent_active( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + agent = _workflow_only_agent() + app = App( + id="app-1", + tenant_id="tenant-1", + name="Workflow", + mode=AppMode.WORKFLOW, + status=AppStatus.NORMAL, + enable_site=True, + enable_api=True, + ) + workflow = Workflow.new( + tenant_id="tenant-1", + app_id=app.id, + type=WorkflowType.WORKFLOW.value, + version=Workflow.VERSION_DRAFT, + graph="{}", + features="{}", + created_by="account-1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + binding = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id=app.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=agent.id, + current_snapshot_id="config-1", + node_job_config={}, + ) + sqlite_session.add_all([agent, app, workflow, binding]) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.retirement_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + result = WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=[agent.id], + account_id="account-1", + ) + + assert result == ([], []) + stored_agent = sqlite_session.get(Agent, agent.id) + assert stored_agent is not None + assert stored_agent.status is AgentStatus.ACTIVE + + +@pytest.mark.parametrize( + "sqlite_session", + [(Agent, App, Workflow, WorkflowAgentNodeBinding, AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], + indirect=True, +) +def test_retire_unowned_archives_orphan_and_retires_resources( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + agent = _workflow_only_agent() + home = AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id=agent.id, + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + workspace = AgentWorkspace( + id="workspace-1", + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id="conversation-1", + owner_scope_key="root", + backend_workspace_ref="workspace-ref", + status=AgentWorkingResourceStatus.ACTIVE, + active_guard=1, + ) + binding = AgentWorkspaceBinding( + id="binding-1", + tenant_id="tenant-1", + app_id="app-1", + workspace_id=workspace.id, + agent_id=agent.id, + base_home_snapshot_id=home.id, + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + sqlite_session.add_all([agent, home, workspace, binding]) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.retirement_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + result = WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=[agent.id], + account_id="account-1", + ) + + assert result == ([binding.id], [home.id]) + stored_agent = sqlite_session.get(Agent, agent.id) + stored_binding = sqlite_session.get(AgentWorkspaceBinding, binding.id) + stored_workspace = sqlite_session.get(AgentWorkspace, workspace.id) + stored_home = sqlite_session.get(AgentHomeSnapshot, home.id) + assert stored_agent is not None + assert stored_binding is not None + assert stored_workspace is not None + assert stored_home is not None + assert stored_agent.status is AgentStatus.ARCHIVED + assert stored_binding.status is AgentWorkingResourceStatus.RETIRED + assert stored_workspace.status is AgentWorkingResourceStatus.RETIRED + assert stored_home.status is AgentWorkingResourceStatus.RETIRED diff --git a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py index ef687714717..ca202fa3047 100644 --- a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py +++ b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py @@ -2,9 +2,13 @@ from types import SimpleNamespace from unittest.mock import ANY, Mock import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session -from models.agent import WorkflowAgentBindingType, WorkflowAgentNodeBinding +from models.agent import Agent, WorkflowAgentBindingType, WorkflowAgentNodeBinding from models.agent_config_entities import WorkflowNodeJobConfig +from models.enums import AppStatus +from models.model import App, AppMode from models.workflow import Workflow, WorkflowType from services.agent.dsl_service import AgentDslService from services.agent.workflow_publish_service import WorkflowAgentPublishService, _InlineAgentOwnershipError @@ -63,16 +67,28 @@ def test_inline_binding_from_another_node_is_cloned(monkeypatch) -> None: assert binding.node_job_config.workflow_prompt == "Summarize the input" -def test_restore_replaces_draft_bindings_with_published_bindings() -> None: - existing = WorkflowAgentNodeBinding( +def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> None: + existing_inline = WorkflowAgentNodeBinding( tenant_id="tenant-1", app_id="app-1", workflow_id="draft-workflow", workflow_version=Workflow.VERSION_DRAFT, - node_id="old-node", + node_id="old-inline-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="old-inline-agent", + current_snapshot_id="old-inline-snapshot", + node_job_config={}, + created_by="account-1", + ) + existing_roster = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="draft-workflow", + workflow_version=Workflow.VERSION_DRAFT, + node_id="old-roster-node", binding_type=WorkflowAgentBindingType.ROSTER_AGENT, - agent_id="old-agent", - current_snapshot_id="old-snapshot", + agent_id="old-roster-agent", + current_snapshot_id="old-roster-snapshot", node_job_config={}, created_by="account-1", ) @@ -89,16 +105,21 @@ def test_restore_replaces_draft_bindings_with_published_bindings() -> None: created_by="account-1", ) session = Mock() - session.scalars.side_effect = [SimpleNamespace(all=lambda: [existing]), SimpleNamespace(all=lambda: [source])] - - WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( + session.scalars.side_effect = [ + SimpleNamespace(all=lambda: [existing_inline, existing_roster]), + SimpleNamespace(all=lambda: [source]), + ] + retirement_candidates = WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( session=session, source_workflow=_workflow(workflow_id="published-workflow", version="2026-07-13 00:00:00"), draft_workflow=_workflow(workflow_id="draft-workflow"), account_id="account-2", ) - session.delete.assert_called_once_with(existing) + assert {item.args[0].agent_id for item in session.delete.call_args_list} == { + "old-inline-agent", + "old-roster-agent", + } restored = session.add.call_args.args[0] assert isinstance(restored, WorkflowAgentNodeBinding) assert restored.workflow_id == "draft-workflow" @@ -107,6 +128,90 @@ def test_restore_replaces_draft_bindings_with_published_bindings() -> None: assert restored.current_snapshot_id == "published-snapshot" assert restored.node_job_config.workflow_prompt == "Use the roster agent" session.flush.assert_called_once() + assert retirement_candidates == {"old-inline-agent"} + + +@pytest.mark.parametrize( + "sqlite_session", + [(App, Agent, WorkflowAgentNodeBinding)], + indirect=True, +) +def test_publish_binding_replacement_returns_only_previous_inline_agent( + sqlite_session: Session, +) -> None: + draft_workflow = _workflow() + draft_workflow.graph = '{"nodes":[{"id":"agent-node","data":{"type":"agent","version":"2"}}],"edges":[]}' + published_workflow = _workflow(workflow_id="published-new", version="published-new") + app = App( + id="app-1", + tenant_id="tenant-1", + name="Workflow", + description="", + mode=AppMode.WORKFLOW, + workflow_id="published-previous", + status=AppStatus.NORMAL, + enable_site=False, + enable_api=False, + api_rpm=0, + api_rph=0, + ) + previous_inline_binding = WorkflowAgentNodeBinding( + id="previous-inline-binding", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="published-previous", + workflow_version="published-previous", + node_id="previous-inline-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="previous-inline-agent", + current_snapshot_id="previous-inline-snapshot", + node_job_config={}, + created_by="account-1", + ) + previous_roster_binding = WorkflowAgentNodeBinding( + id="previous-roster-binding", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="published-previous", + workflow_version="published-previous", + node_id="previous-roster-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="previous-roster-agent", + current_snapshot_id="previous-roster-snapshot", + node_job_config={}, + created_by="account-1", + ) + draft_binding = WorkflowAgentNodeBinding( + id="draft-binding", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version=Workflow.VERSION_DRAFT, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="draft-inline-agent", + current_snapshot_id="draft-inline-snapshot", + node_job_config={"workflow_prompt": "work"}, + created_by="account-1", + ) + sqlite_session.add_all([app, previous_inline_binding, previous_roster_binding, draft_binding]) + sqlite_session.commit() + retirement_candidates = WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + session=sqlite_session, + draft_workflow=draft_workflow, + published_workflow=published_workflow, + ) + + assert retirement_candidates == {"previous-inline-agent"} + copied = sqlite_session.scalar( + select(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.workflow_id == "published-new", + WorkflowAgentNodeBinding.workflow_version == "published-new", + ) + ) + assert copied is not None + assert copied.agent_id == "draft-inline-agent" + assert copied.current_snapshot_id == "draft-inline-snapshot" def test_inline_binding_reuses_existing_node_owned_agent(monkeypatch) -> None: diff --git a/api/tests/unit_tests/services/agent/test_workspace_service.py b/api/tests/unit_tests/services/agent/test_workspace_service.py new file mode 100644 index 00000000000..1dca4dcd1e4 --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_workspace_service.py @@ -0,0 +1,422 @@ +from contextlib import nullcontext +from datetime import datetime, timedelta +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from models.agent import ( + AgentConfigVersionKind, + AgentHomeSnapshot, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) +from services.agent.workspace_service import AgentWorkspaceService, WorkspaceOwnerScope + + +def _scope() -> WorkspaceOwnerScope: + return WorkspaceOwnerScope( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id="conversation-1", + ) + + +def _backend_client() -> MagicMock: + client = MagicMock() + client.create_execution_binding_sync.return_value = SimpleNamespace( + binding_ref="binding-ref", + workspace_ref="workspace-ref", + ) + return client + + +def _workspace( + *, + workspace_id: str = "workspace-1", + tenant_id: str = "tenant-1", + app_id: str = "app-1", + owner_type: AgentWorkspaceOwnerType = AgentWorkspaceOwnerType.CONVERSATION, + owner_id: str = "conversation-1", + status: AgentWorkingResourceStatus = AgentWorkingResourceStatus.ACTIVE, + updated_at: datetime | None = None, + backend_workspace_ref: str = "workspace-ref", +) -> AgentWorkspace: + return AgentWorkspace( + id=workspace_id, + tenant_id=tenant_id, + app_id=app_id, + owner_type=owner_type, + owner_id=owner_id, + owner_scope_key="root", + backend_workspace_ref=backend_workspace_ref, + status=status, + active_guard=1 if status is AgentWorkingResourceStatus.ACTIVE else None, + updated_at=updated_at, + ) + + +def _binding( + *, + binding_id: str = "binding-1", + tenant_id: str = "tenant-1", + app_id: str = "app-1", + workspace_id: str = "workspace-1", + agent_id: str = "agent-1", + status: AgentWorkingResourceStatus = AgentWorkingResourceStatus.ACTIVE, + config_kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT, + updated_at: datetime | None = None, +) -> AgentWorkspaceBinding: + return AgentWorkspaceBinding( + id=binding_id, + tenant_id=tenant_id, + app_id=app_id, + workspace_id=workspace_id, + agent_id=agent_id, + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=config_kind, + backend_binding_ref=f"{binding_id}-ref", + status=status, + updated_at=updated_at, + ) + + +@pytest.mark.parametrize( + "sqlite_session", + [(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], + indirect=True, +) +def test_create_binding_success_persists_new_workspace_and_binding( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + sqlite_session.add( + AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + ) + sqlite_session.commit() + client = _backend_client() + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + + binding = AgentWorkspaceService.create_binding( + session=sqlite_session, + scope=_scope(), + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + sqlite_session.commit() + + workspace = sqlite_session.scalar(select(AgentWorkspace)) + stored_binding = sqlite_session.get(AgentWorkspaceBinding, binding.id) + assert workspace is not None + assert stored_binding is not None + assert binding.workspace_id == workspace.id + assert stored_binding.backend_binding_ref == "binding-ref" + request = client.create_execution_binding_sync.call_args.args[0] + assert request.existing_workspace_ref is None + assert request.workspace_id == workspace.id + + +@pytest.mark.parametrize( + "sqlite_session", + [(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], + indirect=True, +) +def test_create_second_binding_reuses_existing_workspace( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace() + first_binding = _binding() + home = AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + sqlite_session.add_all([workspace, first_binding, home]) + sqlite_session.commit() + client = _backend_client() + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + + binding = AgentWorkspaceService.create_binding( + session=sqlite_session, + scope=_scope(), + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + sqlite_session.commit() + + assert binding.workspace_id == workspace.id + assert sqlite_session.scalars(select(AgentWorkspace)).all() == [workspace] + assert {row.id for row in sqlite_session.scalars(select(AgentWorkspaceBinding)).all()} == { + first_binding.id, + binding.id, + } + request = client.create_execution_binding_sync.call_args.args[0] + assert request.existing_workspace_ref == "workspace-ref" + assert request.workspace_id == workspace.id + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_get_active_binding_resolves_exact_participant(sqlite_session: Session) -> None: + conversation_workspace = _workspace(workspace_id="workspace-conversation") + build_workspace = _workspace( + workspace_id="workspace-build", + owner_type=AgentWorkspaceOwnerType.BUILD_DRAFT, + owner_id="draft-1", + ) + conversation_binding = _binding( + binding_id="binding-conversation", + workspace_id=conversation_workspace.id, + ) + build_binding = _binding( + binding_id="binding-build", + workspace_id=build_workspace.id, + config_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) + sqlite_session.add_all([conversation_workspace, build_workspace, conversation_binding, build_binding]) + sqlite_session.commit() + + resolved = AgentWorkspaceService.get_active_binding( + session=sqlite_session, + tenant_id="tenant-1", + binding_id=conversation_binding.id, + expected_owner_scope=_scope(), + ) + + assert resolved is not None + assert resolved.id == conversation_binding.id + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_get_active_binding_rejects_wrong_owner(sqlite_session: Session) -> None: + build_workspace = _workspace( + workspace_id="workspace-build", + owner_type=AgentWorkspaceOwnerType.BUILD_DRAFT, + owner_id="draft-1", + ) + build_binding = _binding( + binding_id="binding-build", + workspace_id=build_workspace.id, + config_kind=AgentConfigVersionKind.BUILD_DRAFT, + ) + sqlite_session.add_all([build_workspace, build_binding]) + sqlite_session.commit() + + resolved = AgentWorkspaceService.get_active_binding( + session=sqlite_session, + tenant_id="tenant-1", + binding_id=build_binding.id, + expected_owner_scope=_scope(), + ) + + assert resolved is None + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_retire_non_final_binding_keeps_workspace_active(sqlite_session: Session) -> None: + binding = _binding() + other_binding = _binding(binding_id="binding-2", agent_id="agent-2") + workspace = _workspace() + sqlite_session.add_all([workspace, binding, other_binding]) + sqlite_session.commit() + + retired_id = AgentWorkspaceService.retire_binding( + session=sqlite_session, + tenant_id="tenant-1", + binding_id=binding.id, + ) + + assert retired_id == binding.id + assert binding.status is AgentWorkingResourceStatus.RETIRED + assert binding.retired_at is not None + assert workspace.status is AgentWorkingResourceStatus.ACTIVE + assert other_binding.status is AgentWorkingResourceStatus.ACTIVE + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_retire_final_binding_retires_workspace(sqlite_session: Session) -> None: + binding = _binding() + workspace = _workspace() + sqlite_session.add_all([workspace, binding]) + sqlite_session.commit() + + AgentWorkspaceService.retire_binding(session=sqlite_session, tenant_id="tenant-1", binding_id=binding.id) + + assert binding.status is AgentWorkingResourceStatus.RETIRED + assert workspace.status is AgentWorkingResourceStatus.RETIRED + assert workspace.active_guard is None + assert workspace.retired_at == binding.retired_at + + +def test_retire_workspace_retires_all_active_bindings() -> None: + workspace = _workspace() + bindings = [_binding(), _binding(binding_id="binding-2", agent_id="agent-2")] + session = MagicMock() + session.scalar.return_value = workspace + session.scalars.return_value.all.return_value = bindings + + retired_id = AgentWorkspaceService.retire_workspace( + session=session, + tenant_id="tenant-1", + workspace_id=workspace.id, + ) + + assert retired_id == workspace.id + assert workspace.status is AgentWorkingResourceStatus.RETIRED + assert all(binding.status is AgentWorkingResourceStatus.RETIRED for binding in bindings) + assert all(binding.retired_at == workspace.retired_at for binding in bindings) + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_retire_all_for_app_retires_only_active_workspaces_for_that_app(sqlite_session: Session) -> None: + active = _workspace(workspace_id="workspace-active", owner_id="conversation-active") + already_retired = _workspace( + workspace_id="workspace-retired", + owner_id="conversation-retired", + status=AgentWorkingResourceStatus.RETIRED, + ) + other_app = _workspace( + workspace_id="workspace-other", + app_id="app-2", + owner_id="conversation-other", + ) + active_binding = _binding(binding_id="binding-active", workspace_id=active.id) + other_binding = _binding( + binding_id="binding-other", + app_id="app-2", + workspace_id=other_app.id, + ) + sqlite_session.add_all([active, already_retired, other_app, active_binding, other_binding]) + sqlite_session.commit() + + retired_ids = AgentWorkspaceService.retire_all_for_app( + session=sqlite_session, + tenant_id="tenant-1", + app_id="app-1", + ) + + assert retired_ids == [active.id] + assert active.status is AgentWorkingResourceStatus.RETIRED + assert active_binding.status is AgentWorkingResourceStatus.RETIRED + assert already_retired.status is AgentWorkingResourceStatus.RETIRED + assert other_app.status is AgentWorkingResourceStatus.ACTIVE + assert other_binding.status is AgentWorkingResourceStatus.ACTIVE + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_collect_binding_without_retired_workspace_destroys_binding_only( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + binding = _binding(status=AgentWorkingResourceStatus.RETIRED) + workspace = _workspace() + sqlite_session.add_all([workspace, binding]) + sqlite_session.commit() + client = MagicMock() + monkeypatch.setattr( + "services.agent.workspace_service.session_factory.create_session", lambda: nullcontext(sqlite_session) + ) + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + + AgentWorkspaceService.collect_retired_binding(tenant_id="tenant-1", binding_id=binding.id) + + request = client.destroy_execution_binding_sync.call_args.args[0] + assert request.binding_ref == binding.backend_binding_ref + assert request.destroy_workspace is False + assert request.workspace_ref is None + assert sqlite_session.get(AgentWorkspaceBinding, binding.id) is None + assert sqlite_session.get(AgentWorkspace, workspace.id) is not None + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_collect_workspace_destroys_workspace_then_remaining_bindings( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace(status=AgentWorkingResourceStatus.RETIRED) + anchor = _binding(status=AgentWorkingResourceStatus.RETIRED) + remaining = _binding( + binding_id="binding-2", + agent_id="agent-2", + status=AgentWorkingResourceStatus.RETIRED, + ) + anchor.created_at = datetime(2026, 7, 23, 10) + remaining.created_at = anchor.created_at + timedelta(minutes=1) + sqlite_session.add_all([workspace, anchor, remaining]) + sqlite_session.commit() + client = MagicMock() + monkeypatch.setattr( + "services.agent.workspace_service.session_factory.create_session", lambda: nullcontext(sqlite_session) + ) + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + + AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id=workspace.id) + + requests = [call.args[0] for call in client.destroy_execution_binding_sync.call_args_list] + assert len(requests) == 2 + assert requests[0].binding_ref == anchor.backend_binding_ref + assert requests[0].workspace_ref == workspace.backend_workspace_ref + assert requests[0].destroy_workspace is True + assert requests[1].binding_ref == remaining.backend_binding_ref + assert requests[1].destroy_workspace is False + assert sqlite_session.get(AgentWorkspace, workspace.id) is None + assert sqlite_session.get(AgentWorkspaceBinding, anchor.id) is None + assert sqlite_session.get(AgentWorkspaceBinding, remaining.id) is None + + +def test_binding_collection_database_failure_is_best_effort(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + session.scalar.side_effect = RuntimeError("database unavailable") + log_exception = MagicMock() + monkeypatch.setattr("services.agent.workspace_service.session_factory.create_session", lambda: context) + monkeypatch.setattr("services.agent.workspace_service.logger.exception", log_exception) + + AgentWorkspaceService.collect_retired_binding(tenant_id="tenant-1", binding_id="binding-1") + + session.scalar.assert_called_once() + log_exception.assert_called_once_with( + "Failed to collect retired Agent Workspace Binding", + extra={"tenant_id": "tenant-1", "binding_id": "binding-1"}, + ) + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_workspace_collection_final_delete_failure_is_best_effort( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace(status=AgentWorkingResourceStatus.RETIRED) + anchor = _binding(status=AgentWorkingResourceStatus.RETIRED) + sqlite_session.add_all([workspace, anchor]) + sqlite_session.commit() + commit = MagicMock(side_effect=RuntimeError("database unavailable")) + client = MagicMock() + log_exception = MagicMock() + monkeypatch.setattr( + "services.agent.workspace_service.session_factory.create_session", lambda: nullcontext(sqlite_session) + ) + monkeypatch.setattr(sqlite_session, "commit", commit) + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + monkeypatch.setattr("services.agent.workspace_service.logger.exception", log_exception) + + AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id="workspace-1") + + client.destroy_execution_binding_sync.assert_called_once() + log_exception.assert_called_once_with( + "Failed to collect retired Agent Workspace", + extra={"tenant_id": "tenant-1", "workspace_id": "workspace-1"}, + ) diff --git a/api/tests/unit_tests/services/test_agent_app_sandbox_service.py b/api/tests/unit_tests/services/test_agent_app_sandbox_service.py index f3052a4ab49..292b0aa8372 100644 --- a/api/tests/unit_tests/services/test_agent_app_sandbox_service.py +++ b/api/tests/unit_tests/services/test_agent_app_sandbox_service.py @@ -1,470 +1,239 @@ -"""Unit tests for the Agent App / workflow sandbox services.""" - -from __future__ import annotations - -from collections.abc import Generator +from contextlib import nullcontext from datetime import datetime +from types import SimpleNamespace +from unittest.mock import MagicMock import pytest -from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from dify_agent.protocol import RuntimeLayerSpec, SandboxListResponse, SandboxReadResponse, SandboxUploadResponse -from sqlalchemy import delete +from dify_agent.protocol import WorkspaceListResponse, WorkspaceReadResponse +from sqlalchemy.orm import Session -from core.app.apps.agent_app.session_store import AgentAppSessionScope, StoredAgentAppSession -from core.db.session_factory import session_factory -from models.agent import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus +from models.agent import ( + AgentConfigVersionKind, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, +) +from models.enums import ConversationFromSource +from models.model import App, AppMode, Conversation, IconType +from services.agent.workspace_service import AgentWorkspaceService from services.agent_app_sandbox_service import ( AgentAppSandboxService, AgentSandboxInspectorError, - AgentSandboxUploadDownload, WorkflowAgentSandboxService, - _default_client_factory, - _upload_download_response, ) -def _snapshot( - *, - session_id: str = "abc1234", - shell_runtime_state: dict[str, str] | None = None, -) -> CompositorSessionSnapshot: - runtime_state = ( - {"session_id": session_id, "workspace_cwd": f"~/workspace/{session_id}"} - if shell_runtime_state is None - else shell_runtime_state - ) - return CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot(name="execution_context", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), - LayerSessionSnapshot( - name="shell", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state=runtime_state, - ), - ] - ) - - -def _runtime_layer_specs() -> list[RuntimeLayerSpec]: - return [ - RuntimeLayerSpec(name="execution_context", type="dify.execution_context", config={"tenant_id": "tenant-1"}), - RuntimeLayerSpec(name="shell", type="dify.shell", deps={"execution_context": "execution_context"}, config={}), - ] - - -class FakeStore: - def __init__(self, session: StoredAgentAppSession | None) -> None: - self.session = session - self.scope: tuple[str, str, str] | None = None - - def load_active_session_for_conversation(self, *, tenant_id: str, app_id: str, conversation_id: str): - self.scope = (tenant_id, app_id, conversation_id) - return self.session - - -class FakeClient: - def __init__(self) -> None: - self.calls: list[tuple[str, str]] = [] - self.locators: list[object] = [] - - def list_sandbox_files_sync(self, locator, path: str) -> SandboxListResponse: - self.locators.append(locator) - self.calls.append(("list", path)) - return SandboxListResponse(path=path, entries=[], truncated=False) - - def read_sandbox_file_sync(self, locator, path: str, max_bytes: int = 262144) -> SandboxReadResponse: - del max_bytes - self.locators.append(locator) - self.calls.append(("read", path)) - return SandboxReadResponse(path=path, size=5, truncated=False, binary=False, text="hello") - - def upload_sandbox_file_sync(self, locator, path: str) -> SandboxUploadResponse: - self.locators.append(locator) - self.calls.append(("upload", path)) - return SandboxUploadResponse( - path=path, - file={ - "transfer_method": "tool_file", - "reference": "dify-file-ref:file-1", - "download_url": "https://files.example/report.txt?token=1", - }, - ) - - -def _stored_session( - *, - session_snapshot: CompositorSessionSnapshot | None = None, -) -> StoredAgentAppSession: - return StoredAgentAppSession( - scope=AgentAppSessionScope( +def _add_normal_conversation(session: Session, *, binding_id: str) -> Conversation: + session.add( + App( + id="app-1", tenant_id="tenant-1", - app_id="app-1", - conversation_id="conv-1", - agent_id="agent-1", - agent_config_snapshot_id="snapshot-1", - ), - session_snapshot=_snapshot() if session_snapshot is None else session_snapshot, - backend_run_id="run-1", - runtime_layer_specs=_runtime_layer_specs(), - ) - - -def test_agent_app_sandbox_service_get_info_returns_metadata() -> None: - store = FakeStore(_stored_session()) - client = FakeClient() - service = AgentAppSandboxService(session_store=store, client_factory=lambda: client) # type: ignore[arg-type] - - result = service.get_info(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1") - - assert result.session_id == "abc1234" - assert result.workspace_cwd == "~/workspace/abc1234" - assert client.calls == [] - assert store.scope == ("tenant-1", "app-1", "conv-1") - - -def test_agent_app_sandbox_service_builds_locator_and_proxies() -> None: - store = FakeStore(_stored_session()) - client = FakeClient() - service = AgentAppSandboxService(session_store=store, client_factory=lambda: client) # type: ignore[arg-type] - - result = service.list_files(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1", path=".") - - assert result.path == "." - assert client.calls == [("list", ".")] - assert store.scope == ("tenant-1", "app-1", "conv-1") - - -def test_agent_app_sandbox_service_upload_returns_download_url(monkeypatch: pytest.MonkeyPatch) -> None: - store = FakeStore(_stored_session()) - client = FakeClient() - captured: dict[str, object] = {} - - def fake_upload_download_response(*, tenant_id: str, file_mapping: dict[str, object]) -> AgentSandboxUploadDownload: - captured["tenant_id"] = tenant_id - captured["file_mapping"] = file_mapping - return AgentSandboxUploadDownload(url="https://files.example/report.txt?token=1&as_attachment=true") - - monkeypatch.setattr("services.agent_app_sandbox_service._upload_download_response", fake_upload_download_response) - service = AgentAppSandboxService(session_store=store, client_factory=lambda: client) # type: ignore[arg-type] - - result = service.upload_file(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1", path="report.txt") - - assert result.url == "https://files.example/report.txt?token=1&as_attachment=true" - assert client.calls == [("upload", "report.txt")] - assert store.scope == ("tenant-1", "app-1", "conv-1") - assert captured == { - "tenant_id": "tenant-1", - "file_mapping": { - "transfer_method": "tool_file", - "reference": "dify-file-ref:file-1", - "download_url": "https://files.example/report.txt?token=1", - }, - } - - -def test_agent_app_sandbox_service_raises_when_no_active_session() -> None: - service = AgentAppSandboxService(session_store=FakeStore(None), client_factory=lambda: FakeClient()) # type: ignore[arg-type] - - with pytest.raises(AgentSandboxInspectorError) as exc_info: - service.get_info(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1") - - assert exc_info.value.code == "no_active_session" - assert exc_info.value.status_code == 404 - - -def test_agent_app_sandbox_service_raises_when_shell_workspace_metadata_missing() -> None: - broken_session = _stored_session(session_snapshot=_snapshot(shell_runtime_state={"session_id": "abc1234"})) - service = AgentAppSandboxService(session_store=FakeStore(broken_session), client_factory=lambda: FakeClient()) # type: ignore[arg-type] - - with pytest.raises(AgentSandboxInspectorError) as exc_info: - service.get_info(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1") - - assert exc_info.value.code == "no_sandbox" - assert exc_info.value.status_code == 404 - - -def test_default_client_factory_requires_agent_backend_base_url(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr("services.agent_app_sandbox_service.dify_config.AGENT_BACKEND_BASE_URL", "") - - with pytest.raises(AgentSandboxInspectorError) as exc_info: - _default_client_factory() - - assert exc_info.value.code == "inspector_unavailable" - assert exc_info.value.status_code == 503 - - -@pytest.fixture -def _runtime_session_table() -> Generator[None, None, None]: - engine = session_factory.get_session_maker().kw["bind"] - AgentRuntimeSession.__table__.create(bind=engine, checkfirst=True) - yield - with session_factory.create_session() as session: - session.execute(delete(AgentRuntimeSession)) - session.commit() - AgentRuntimeSession.__table__.drop(bind=engine, checkfirst=True) - - -def _insert_workflow_session( - *, - runtime_layer_specs: str | None = None, - workflow_run_id: str = "run-1", - node_id: str = "node-1", - node_execution_id: str = "node-exec-1", - binding_id: str = "binding-1", - backend_run_id: str = "backend-run-1", - updated_at: datetime | None = None, - session_id: str = "abc1234", -) -> None: - default_runtime_layer_specs = ( - '[{"name":"execution_context","type":"dify.execution_context","config":{"tenant_id":"tenant-1"}},' - '{"name":"shell","type":"dify.shell","deps":{"execution_context":"execution_context"},"config":{}}]' - ) - with session_factory.create_session() as session: - row = AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.WORKFLOW_RUN, - workflow_id="workflow-1", - workflow_run_id=workflow_run_id, - node_id=node_id, - node_execution_id=node_execution_id, - binding_id=binding_id, - agent_id="agent-1", - agent_config_snapshot_id="snapshot-1", - backend_run_id=backend_run_id, - session_snapshot=_snapshot(session_id=session_id).model_dump_json(), - composition_layer_specs=runtime_layer_specs or default_runtime_layer_specs, - status=AgentRuntimeSessionStatus.ACTIVE, + name="Agent App", + description="", + mode=AppMode.AGENT, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=False, + enable_api=False, + max_active_requests=0, ) - if updated_at is not None: - row.updated_at = updated_at - session.add(row) - session.commit() + ) + conversation = Conversation( + id="conversation-1", + app_id="app-1", + mode=AppMode.AGENT, + name="Conversation", + from_source=ConversationFromSource.CONSOLE, + from_account_id="account-1", + is_deleted=False, + agent_workspace_binding_id=binding_id, + ) + conversation._inputs = {} + session.add(conversation) + return conversation -@pytest.mark.usefixtures("_runtime_session_table") -def test_workflow_sandbox_service_resolves_locator_and_returns_download_url( - monkeypatch: pytest.MonkeyPatch, -) -> None: - _insert_workflow_session() - client = FakeClient() - captured: dict[str, object] = {} - - def fake_upload_download_response(*, tenant_id: str, file_mapping: dict[str, object]) -> AgentSandboxUploadDownload: - captured["tenant_id"] = tenant_id - captured["file_mapping"] = file_mapping - return AgentSandboxUploadDownload(url="https://files.example/report.txt?token=1&as_attachment=true") - - monkeypatch.setattr("services.agent_app_sandbox_service._upload_download_response", fake_upload_download_response) - service = WorkflowAgentSandboxService(client_factory=lambda: client) # type: ignore[arg-type] - - result = service.upload_file( +def _add_conversation_bindings(session: Session) -> tuple[AgentWorkspaceBinding, AgentWorkspaceBinding]: + workspace = AgentWorkspace( + id="workspace-1", tenant_id="tenant-1", app_id="app-1", - workflow_run_id="run-1", - node_id="node-1", - node_execution_id="node-exec-1", - path="report.txt", - session=session_factory.create_session(), + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id="conversation-1", + owner_scope_key="root", + backend_workspace_ref="workspace-ref", + status=AgentWorkingResourceStatus.ACTIVE, + active_guard=1, ) - - assert result.url == "https://files.example/report.txt?token=1&as_attachment=true" - assert client.calls == [("upload", "report.txt")] - assert captured == { - "tenant_id": "tenant-1", - "file_mapping": { - "transfer_method": "tool_file", - "reference": "dify-file-ref:file-1", - "download_url": "https://files.example/report.txt?token=1", - }, - } - - -def test_upload_download_response_resolves_signed_external_url(monkeypatch: pytest.MonkeyPatch) -> None: - built_file = object() - built_with: dict[str, object] = {} - - def fake_build_from_mapping(*, mapping: dict[str, object], tenant_id: str, access_controller: object) -> object: - built_with["mapping"] = mapping - built_with["tenant_id"] = tenant_id - built_with["access_controller"] = access_controller - return built_file - - class FakeRuntime: - def __init__(self, *, file_access_controller: object) -> None: - self.file_access_controller = file_access_controller - - def resolve_file_url(self, *, file: object, for_external: bool) -> str: - assert file is built_file - assert for_external is True - return "https://files.example/files/tools/tool-file.txt?timestamp=1&nonce=2&sign=3" - - monkeypatch.setattr("services.agent_app_sandbox_service.file_factory.build_from_mapping", fake_build_from_mapping) - monkeypatch.setattr("services.agent_app_sandbox_service.DifyWorkflowFileRuntime", FakeRuntime) - - result = _upload_download_response( - tenant_id="tenant-1", - file_mapping={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, - ) - - assert result.url == ( - "https://files.example/files/tools/tool-file.txt?timestamp=1&nonce=2&sign=3&as_attachment=true" - ) - assert built_with["mapping"] == {"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"} - assert built_with["tenant_id"] == "tenant-1" - assert built_with["access_controller"] is not None - - -def test_upload_download_response_maps_resolution_failure_to_inspector_error( - monkeypatch: pytest.MonkeyPatch, -) -> None: - def fake_build_from_mapping(*, mapping: dict[str, object], tenant_id: str, access_controller: object) -> object: - del mapping, tenant_id, access_controller - raise ValueError("missing tool file") - - monkeypatch.setattr("services.agent_app_sandbox_service.file_factory.build_from_mapping", fake_build_from_mapping) - - with pytest.raises(AgentSandboxInspectorError) as exc_info: - _upload_download_response( - tenant_id="tenant-1", - file_mapping={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, - ) - - assert exc_info.value.code == "sandbox_upload_download_unavailable" - assert exc_info.value.status_code == 502 - - -def test_upload_download_response_maps_missing_url_to_inspector_error( - monkeypatch: pytest.MonkeyPatch, -) -> None: - built_file = object() - - def fake_build_from_mapping(*, mapping: dict[str, object], tenant_id: str, access_controller: object) -> object: - del mapping, tenant_id, access_controller - return built_file - - class FakeRuntime: - def __init__(self, *, file_access_controller: object) -> None: - self.file_access_controller = file_access_controller - - def resolve_file_url(self, *, file: object, for_external: bool) -> None: - assert file is built_file - assert for_external is True - - monkeypatch.setattr("services.agent_app_sandbox_service.file_factory.build_from_mapping", fake_build_from_mapping) - monkeypatch.setattr("services.agent_app_sandbox_service.DifyWorkflowFileRuntime", FakeRuntime) - - with pytest.raises(AgentSandboxInspectorError) as exc_info: - _upload_download_response( - tenant_id="tenant-1", - file_mapping={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, - ) - - assert exc_info.value.code == "sandbox_upload_download_unavailable" - assert exc_info.value.status_code == 502 - - -@pytest.mark.usefixtures("_runtime_session_table") -def test_workflow_sandbox_service_filters_by_node_execution_id() -> None: - _insert_workflow_session( - node_execution_id="node-exec-1", - binding_id="binding-1", - backend_run_id="run-a", - session_id="abc1234", - ) - _insert_workflow_session( - node_execution_id="node-exec-2", - binding_id="binding-2", - backend_run_id="run-b", - session_id="def5678", - ) - client = FakeClient() - service = WorkflowAgentSandboxService(client_factory=lambda: client) # type: ignore[arg-type] - - result = service.read_file( + expected = AgentWorkspaceBinding( + id="binding-expected", tenant_id="tenant-1", app_id="app-1", - workflow_run_id="run-1", - node_id="node-1", - node_execution_id="node-exec-2", - path="out.txt", - session=session_factory.create_session(), + workspace_id=workspace.id, + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-expected-ref", + status=AgentWorkingResourceStatus.ACTIVE, + updated_at=datetime(2026, 7, 23, 10), ) - - assert result.text == "hello" - assert client.calls == [("read", "out.txt")] - assert client.locators[0].session_snapshot.layers[1].runtime_state["session_id"] == "def5678" - - -@pytest.mark.usefixtures("_runtime_session_table") -def test_workflow_sandbox_service_uses_latest_active_session_when_execution_id_omitted() -> None: - _insert_workflow_session( - node_execution_id="node-exec-1", - binding_id="binding-1", - backend_run_id="run-older", - updated_at=datetime(2026, 1, 1, 0, 0, 0), - session_id="abc1234", - ) - _insert_workflow_session( - node_execution_id="node-exec-2", - binding_id="binding-2", - backend_run_id="run-newer", - updated_at=datetime(2026, 1, 1, 0, 0, 1), - session_id="def5678", - ) - client = FakeClient() - service = WorkflowAgentSandboxService(client_factory=lambda: client) # type: ignore[arg-type] - - result = service.list_files( + other = AgentWorkspaceBinding( + id="binding-other", tenant_id="tenant-1", app_id="app-1", - workflow_run_id="run-1", - node_id="node-1", - node_execution_id=None, + workspace_id=workspace.id, + agent_id="agent-1", + base_home_snapshot_id="home-1", + agent_config_version_id="config-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-other-ref", + status=AgentWorkingResourceStatus.ACTIVE, + updated_at=datetime(2026, 7, 23, 11), + ) + session.add_all([workspace, expected, other]) + return expected, other + + +def _use_session(monkeypatch: pytest.MonkeyPatch, session: Session) -> None: + monkeypatch.setattr( + "services.agent_app_sandbox_service.session_factory.create_session", + lambda: nullcontext(session), + ) + + +@pytest.mark.parametrize( + "sqlite_session", + [(AgentWorkspace, AgentWorkspaceBinding, App, Conversation)], + indirect=True, +) +def test_agent_app_file_browsing_uses_conversation_pointer( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + expected, _ = _add_conversation_bindings(sqlite_session) + _add_normal_conversation(sqlite_session, binding_id=expected.id) + sqlite_session.commit() + _use_session(monkeypatch, sqlite_session) + client = MagicMock() + response = WorkspaceListResponse(path=".", entries=[], truncated=False) + client.list_workspace_files_sync.return_value = response + + result = AgentAppSandboxService(client_factory=lambda: nullcontext(client)).list_files( + tenant_id="tenant-1", + app_id="app-1", + agent_id="agent-1", + caller_type="conversation", + caller_id="conversation-1", + account_id="account-1", path=".", - session=session_factory.create_session(), ) - assert result.path == "." - assert client.calls == [("list", ".")] - assert client.locators[0].session_snapshot.layers[1].runtime_state["session_id"] == "def5678" + assert result is response + client.list_workspace_files_sync.assert_called_once_with(expected.backend_binding_ref, ".") -@pytest.mark.usefixtures("_runtime_session_table") -def test_workflow_sandbox_service_raises_when_no_active_session() -> None: - service = WorkflowAgentSandboxService(client_factory=lambda: FakeClient()) # type: ignore[arg-type] +@pytest.mark.parametrize( + "sqlite_session", + [(AgentWorkspace, AgentWorkspaceBinding, App, Conversation)], + indirect=True, +) +def test_agent_app_file_browsing_rejects_other_account( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + expected, _ = _add_conversation_bindings(sqlite_session) + _add_normal_conversation(sqlite_session, binding_id=expected.id) + sqlite_session.commit() + _use_session(monkeypatch, sqlite_session) + client = MagicMock() with pytest.raises(AgentSandboxInspectorError) as exc_info: - service.list_files( + AgentAppSandboxService(client_factory=lambda: nullcontext(client)).list_files( tenant_id="tenant-1", app_id="app-1", - workflow_run_id="run-1", - node_id="node-1", - node_execution_id=None, + agent_id="agent-1", + caller_type="conversation", + caller_id="conversation-1", + account_id="account-2", path=".", - session=session_factory.create_session(), ) - assert exc_info.value.code == "no_active_session" - assert exc_info.value.status_code == 404 + assert exc_info.value.code == "no_active_binding" + client.list_workspace_files_sync.assert_not_called() -@pytest.mark.usefixtures("_runtime_session_table") -def test_workflow_sandbox_service_raises_when_runtime_specs_missing() -> None: - _insert_workflow_session(runtime_layer_specs="[]") - service = WorkflowAgentSandboxService(client_factory=lambda: FakeClient()) # type: ignore[arg-type] +@pytest.mark.parametrize( + ("parent_app_id", "backing_app_id", "runtime_app_id"), + [ + ("app-1", None, "app-1"), + ("workflow-app-1", "runtime-app-1", "runtime-app-1"), + ], +) +def test_agent_app_file_browsing_uses_build_draft_caller( + monkeypatch: pytest.MonkeyPatch, + parent_app_id: str, + backing_app_id: str | None, + runtime_app_id: str, +) -> None: + session = MagicMock() + session.scalar.side_effect = [ + SimpleNamespace(app_id=parent_app_id, backing_app_id=backing_app_id), + SimpleNamespace(agent_workspace_binding_id="binding-build"), + ] + _use_session(monkeypatch, session) + binding = SimpleNamespace( + agent_id="agent-1", + backend_binding_ref="binding-build-ref", + ) + get_binding = MagicMock(return_value=binding) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding) + client = MagicMock() + response = WorkspaceListResponse(path=".", entries=[], truncated=False) + client.list_workspace_files_sync.return_value = response - with pytest.raises(AgentSandboxInspectorError) as exc_info: - service.list_files( - tenant_id="tenant-1", - app_id="app-1", - workflow_run_id="run-1", - node_id="node-1", - node_execution_id=None, - path=".", - session=session_factory.create_session(), - ) + result = AgentAppSandboxService(client_factory=lambda: nullcontext(client)).list_files( + tenant_id="tenant-1", + app_id=runtime_app_id, + agent_id="agent-1", + caller_type="build_draft", + caller_id="build-1", + account_id="account-1", + path=".", + ) - assert exc_info.value.code == "no_sandbox" + assert result is response + owner_scope = get_binding.call_args.kwargs["expected_owner_scope"] + assert owner_scope.app_id == runtime_app_id + assert owner_scope.owner_type is AgentWorkspaceOwnerType.BUILD_DRAFT + assert owner_scope.owner_id == "build-1" + client.list_workspace_files_sync.assert_called_once_with("binding-build-ref", ".") + + +def test_workflow_file_access_uses_node_execution_pointer(monkeypatch: pytest.MonkeyPatch) -> None: + execution = SimpleNamespace( + agent_workspace_binding_id="binding-workflow", + process_data_dict={"workflow_agent_binding_id": "workflow-binding-1"}, + ) + session = MagicMock() + session.scalar.return_value = execution + binding = SimpleNamespace(backend_binding_ref="binding-workflow-ref") + get_binding = MagicMock(return_value=binding) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding) + client = MagicMock() + response = WorkspaceReadResponse(path="report.txt", size=2, truncated=False, binary=False, text="ok") + client.read_workspace_file_sync.return_value = response + + result = WorkflowAgentSandboxService(client_factory=lambda: nullcontext(client)).read_file( + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + node_id="node-1", + node_execution_id="execution-1", + path="report.txt", + session=session, + ) + + assert result is response + assert get_binding.call_args.kwargs["binding_id"] == "binding-workflow" + client.read_workspace_file_sync.assert_called_once_with("binding-workflow-ref", "report.txt") diff --git a/api/tests/unit_tests/services/test_app_service.py b/api/tests/unit_tests/services/test_app_service.py index 1e2d033e6ef..c7911534d6c 100644 --- a/api/tests/unit_tests/services/test_app_service.py +++ b/api/tests/unit_tests/services/test_app_service.py @@ -626,7 +626,8 @@ class TestAgentAppType: from services.app_service import AppService app = SimpleNamespace(id="app-1", tenant_id="tenant-1", mode=AppMode.AGENT) - backing_agent = SimpleNamespace(status=AgentStatus.ACTIVE, archived_by=None, archived_at=None) + backing_agent = SimpleNamespace(id="agent-1", status=AgentStatus.ACTIVE, archived_by=None, archived_at=None) + events: list[str] = [] with ( patch("services.app_service.db") as mock_db, @@ -636,11 +637,81 @@ class TestAgentAppType: patch("services.app_service.FeatureService"), patch("services.app_service.dify_config"), patch("services.app_service.remove_app_and_related_data_task"), + patch( + "services.app_service.AgentHomeSnapshotService.retire_all_for_agent", + return_value=["home-1"], + ) as mock_retire_homes, + patch( + "services.app_service.AgentWorkspaceService.retire_all_for_app", + side_effect=lambda **_kwargs: events.append("retire-app-workspaces") or ["workspace-1"], + ) as mock_retire_workspaces, + patch( + "services.app_service.WorkflowAgentRetirementService.retire_unowned", + side_effect=lambda **_kwargs: ( + events.append("retire-workflow-agents") or (["workflow-binding-1"], ["workflow-home-1"]) + ), + ) as mock_workflow_retirement, + patch( + "services.app_service.enqueue_agent_resource_collection", + side_effect=lambda **_kwargs: events.append("enqueue"), + ) as mock_enqueue_collection, ): mock_db.session.scalar.return_value = backing_agent + mock_db.session.commit.side_effect = lambda: events.append("commit") + workflow_agents = MagicMock() + workflow_agents.all.return_value = ["workflow-agent-1", "workflow-agent-2"] + bindings = MagicMock() + bindings.all.return_value = [] + mock_db.session.scalars.side_effect = [workflow_agents, bindings] AppService().delete_app(app, session=mock_db.session) # type: ignore[arg-type] + assert events == ["retire-app-workspaces", "commit", "retire-workflow-agents", "enqueue"] assert backing_agent.status == AgentStatus.ARCHIVED assert backing_agent.archived_by == "account-2" assert backing_agent.archived_at is not None mock_db.session.delete.assert_called_once_with(app) + mock_workflow_retirement.assert_called_once_with( + tenant_id="tenant-1", + agent_ids=["workflow-agent-1", "workflow-agent-2"], + account_id="account-2", + ) + mock_retire_workspaces.assert_called_once_with( + session=mock_db.session, + tenant_id="tenant-1", + app_id="app-1", + ) + mock_retire_homes.assert_called_once_with( + session=mock_db.session, + tenant_id="tenant-1", + agent_id="agent-1", + ) + mock_enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + workspace_ids=["workspace-1"], + binding_ids=["workflow-binding-1"], + home_snapshot_ids=["home-1", "workflow-home-1"], + ) + + def test_delete_app_commit_failure_does_not_retire_workflow_agents_or_enqueue(self): + from models.model import AppMode + from services.app_service import AppService + + app = SimpleNamespace(id="app-1", tenant_id="tenant-1", mode=AppMode.WORKFLOW) + with ( + patch("services.app_service.db") as mock_db, + patch("services.app_service.current_user", SimpleNamespace(id="account-2")), + patch("services.app_service.AgentWorkspaceService.retire_all_for_app", return_value=["workspace-1"]), + patch("services.app_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + patch("services.app_service.enqueue_agent_resource_collection") as enqueue_collection, + ): + mock_db.session.scalar.return_value = None + workflow_agents = MagicMock() + workflow_agents.all.return_value = ["workflow-agent-1"] + mock_db.session.scalars.return_value = workflow_agents + mock_db.session.commit.side_effect = RuntimeError("commit failed") + + with pytest.raises(RuntimeError, match="commit failed"): + AppService().delete_app(app, session=mock_db.session) # type: ignore[arg-type] + + retire_unowned.assert_not_called() + enqueue_collection.assert_not_called() diff --git a/api/tests/unit_tests/services/test_conversation_service.py b/api/tests/unit_tests/services/test_conversation_service.py index 90727705061..79c138ed265 100644 --- a/api/tests/unit_tests/services/test_conversation_service.py +++ b/api/tests/unit_tests/services/test_conversation_service.py @@ -8,7 +8,7 @@ in-memory SQLite sessions with persisted ORM rows. """ import json -from unittest.mock import patch +from unittest.mock import MagicMock, Mock, patch import pytest from sqlalchemy import asc, desc @@ -19,6 +19,8 @@ from libs.datetime_utils import naive_utc_now from models import Account, ConversationVariable from models.enums import AppStatus, ConversationFromSource, ConversationStatus from models.model import App, AppMode, Conversation +from services import conversation_service +from services.agent.workspace_service import AgentWorkspaceService from services.conversation_service import ConversationService TENANT_ID = "11111111-1111-1111-1111-111111111111" @@ -148,6 +150,58 @@ class ConversationServiceTestDataFactory: return conversation +def test_delete_retires_then_commits_before_enqueue(monkeypatch: pytest.MonkeyPatch) -> None: + app = ConversationServiceTestDataFactory.create_app() + conversation = ConversationServiceTestDataFactory.create_conversation() + conversation.agent_workspace_binding_id = "conversation-binding-1" + session = MagicMock() + events: list[str] = [] + get_binding = MagicMock(return_value=Mock(id="conversation-binding-1")) + retire_binding = MagicMock(side_effect=lambda **_kwargs: events.append("retire") or "conversation-binding-1") + monkeypatch.setattr(ConversationService, "get_conversation", MagicMock(return_value=conversation)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", retire_binding) + session.commit.side_effect = lambda: events.append("commit") + monkeypatch.setattr( + conversation_service, + "enqueue_agent_resource_collection", + MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), + ) + monkeypatch.setattr(conversation_service.delete_conversation_related_data, "delay", MagicMock()) + + ConversationService.delete(app, conversation.id, None, session=session) + + assert events == ["retire", "commit", "enqueue"] + assert get_binding.call_args.kwargs["binding_id"] == "conversation-binding-1" + assert retire_binding.call_args.kwargs["binding_id"] == "conversation-binding-1" + + +def test_delete_commit_failure_does_not_enqueue(monkeypatch: pytest.MonkeyPatch) -> None: + app = ConversationServiceTestDataFactory.create_app() + conversation = ConversationServiceTestDataFactory.create_conversation() + conversation.agent_workspace_binding_id = "binding-1" + session = MagicMock() + session.commit.side_effect = RuntimeError("commit failed") + monkeypatch.setattr(ConversationService, "get_conversation", MagicMock(return_value=conversation)) + monkeypatch.setattr( + AgentWorkspaceService, + "get_active_binding", + MagicMock(return_value=Mock(id="binding-1")), + ) + monkeypatch.setattr(AgentWorkspaceService, "retire_binding", MagicMock(return_value="binding-1")) + enqueue_collection = MagicMock() + delete_related = MagicMock() + monkeypatch.setattr(conversation_service, "enqueue_agent_resource_collection", enqueue_collection) + monkeypatch.setattr(conversation_service.delete_conversation_related_data, "delay", delete_related) + + with pytest.raises(RuntimeError, match="commit failed"): + ConversationService.delete(app, conversation.id, None, session=session) + + session.rollback.assert_called_once_with() + enqueue_collection.assert_not_called() + delete_related.assert_not_called() + + @pytest.mark.parametrize("sqlite_session", [(Conversation,)], indirect=True) class TestConversationServicePagination: """Test conversation pagination operations.""" diff --git a/api/tests/unit_tests/services/test_snippet_dsl_service.py b/api/tests/unit_tests/services/test_snippet_dsl_service.py index e527e5bd950..6ae44cbd588 100644 --- a/api/tests/unit_tests/services/test_snippet_dsl_service.py +++ b/api/tests/unit_tests/services/test_snippet_dsl_service.py @@ -492,6 +492,7 @@ def test_check_dependencies_returns_generated_dependencies(monkeypatch): def test_create_or_update_snippet_updates_existing_snippet_and_syncs_workflow(monkeypatch): snippet = SimpleNamespace( id="snippet-1", + tenant_id="tenant-1", name="Old", description="Old", type="node", @@ -508,6 +509,14 @@ def test_create_or_update_snippet_updates_existing_snippet_and_syncs_workflow(mo sync_draft_workflow=Mock(), ) monkeypatch.setattr("services.snippet_dsl_service.SnippetService", lambda *_args, **_kwargs: snippet_service) + monkeypatch.setattr( + "services.snippet_dsl_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", + Mock(return_value=set()), + ) + monkeypatch.setattr( + "services.snippet_dsl_service.WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync", + Mock(), + ) result = service._create_or_update_snippet( snippet=snippet, @@ -537,6 +546,14 @@ def test_create_or_update_snippet_creates_new_snippet_and_flushes(monkeypatch): service = SnippetDslService(session=session) snippet_service = SimpleNamespace(get_draft_workflow=Mock(return_value=None), sync_draft_workflow=Mock()) monkeypatch.setattr("services.snippet_dsl_service.SnippetService", lambda *_args, **_kwargs: snippet_service) + monkeypatch.setattr( + "services.snippet_dsl_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", + Mock(return_value=set()), + ) + monkeypatch.setattr( + "services.snippet_dsl_service.WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync", + Mock(), + ) result = service._create_or_update_snippet( snippet=None, diff --git a/api/tests/unit_tests/services/test_snippet_service.py b/api/tests/unit_tests/services/test_snippet_service.py index b3c6cdcbf6c..799fd4a53a0 100644 --- a/api/tests/unit_tests/services/test_snippet_service.py +++ b/api/tests/unit_tests/services/test_snippet_service.py @@ -58,6 +58,7 @@ def test_create_snippet_allows_duplicate_names(monkeypatch: pytest.MonkeyPatch) account = SimpleNamespace(id="account-1") service = SnippetService.__new__(SnippetService) + service._session = None service._session_maker = _session_maker(session) snippet = service.create_snippet( @@ -190,6 +191,7 @@ def test_update_snippet_updates_optional_fields() -> None: def test_sync_draft_workflow_creates_draft_and_updates_input_fields(monkeypatch: pytest.MonkeyPatch) -> None: service = SnippetService.__new__(SnippetService) + service._session = None monkeypatch.setattr(service, "get_draft_workflow", Mock(return_value=None)) session = Mock() session.scalars.return_value.all.return_value = [] @@ -235,6 +237,7 @@ def test_sync_draft_workflow_raises_when_hash_mismatches() -> None: def test_sync_draft_workflow_updates_existing_draft_and_clears_variables(monkeypatch: pytest.MonkeyPatch) -> None: service = SnippetService.__new__(SnippetService) + service._session = None workflow = _create_workflow( workflow_id="workflow-1", version=Workflow.VERSION_DRAFT, @@ -376,6 +379,7 @@ def test_restore_published_snippet_workflow_to_draft_copies_source_snapshot( features={}, ) service = SnippetService.__new__(SnippetService) + service._session = None session = Mock() session.scalars.return_value.all.return_value = [] service._session_maker = _session_maker(session) @@ -431,6 +435,7 @@ def test_restore_published_snippet_workflow_to_draft_adds_new_draft(monkeypatch: features={}, ) service = SnippetService.__new__(SnippetService) + service._session = None session = Mock() session.scalars.return_value.all.return_value = [] service._session_maker = _session_maker(session) @@ -465,6 +470,7 @@ def test_get_published_workflow_by_id_raises_for_draft(monkeypatch: pytest.Monke draft_workflow = SimpleNamespace(version=Workflow.VERSION_DRAFT) session = SimpleNamespace(scalar=Mock(return_value=draft_workflow)) service = SnippetService.__new__(SnippetService) + service._session = None service._session_maker = _session_maker(session) with pytest.raises(IsDraftWorkflowError): @@ -503,8 +509,12 @@ def test_publish_workflow_creates_snapshot_and_updates_snippet(monkeypatch: pyte updated_by=None, ) session = SimpleNamespace(scalar=Mock(return_value=draft_workflow), add=Mock()) + monkeypatch.setattr( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.copy_agent_node_bindings_to_published", + Mock(return_value=set()), + ) - result = service.publish_workflow( + result, retirement_candidates = service.publish_workflow( session=session, snippet=snippet, account=SimpleNamespace(id="account-1"), @@ -516,6 +526,7 @@ def test_publish_workflow_creates_snapshot_and_updates_snippet(monkeypatch: pyte assert snippet.workflow_id == result.id assert snippet.updated_by == "account-1" assert session.add.call_args_list[-1].args == (snippet,) + assert retirement_candidates == set() def test_get_all_published_workflows_returns_empty_without_current_workflow() -> None: diff --git a/api/tests/unit_tests/services/test_workflow_service.py b/api/tests/unit_tests/services/test_workflow_service.py index b2e0e4129c9..9248d829324 100644 --- a/api/tests/unit_tests/services/test_workflow_service.py +++ b/api/tests/unit_tests/services/test_workflow_service.py @@ -832,7 +832,7 @@ class TestWorkflowService: patch("services.workflow_service.app_published_workflow_was_updated"), patch("services.workflow_service.dify_config.BILLING_ENABLED", False), ): - result = workflow_service.publish_workflow( + result, retirement_candidates = workflow_service.publish_workflow( session=sqlite_session, app_model=app, account=account, @@ -845,6 +845,7 @@ class TestWorkflowService: assert result.version != Workflow.VERSION_DRAFT assert result.marked_name == "Version 1" assert result.marked_comment == "Initial release" + assert retirement_candidates == set() def test_publish_workflow_no_draft_raises_error(self, workflow_service: WorkflowService, sqlite_session: Session): """ diff --git a/api/tests/unit_tests/tasks/test_agent_backend_session_cleanup_task.py b/api/tests/unit_tests/tasks/test_agent_backend_session_cleanup_task.py deleted file mode 100644 index f7bc8a4a9ac..00000000000 --- a/api/tests/unit_tests/tasks/test_agent_backend_session_cleanup_task.py +++ /dev/null @@ -1,51 +0,0 @@ -import logging - -from agenton.compositor import CompositorSessionSnapshot - -from clients.agent_backend.session_cleanup import ( - AgentBackendSessionCleanupPayload, - AgentBackendSessionCleanupResult, -) -from tasks import agent_backend_session_cleanup_task as cleanup_task_module - - -def _payload_dict() -> dict[str, object]: - return AgentBackendSessionCleanupPayload( - session_snapshot=CompositorSessionSnapshot(layers=[]), - runtime_layer_specs=[], - metadata={"tenant_id": "tenant-1", "app_id": "app-1"}, - ).model_dump(mode="json") - - -def test_run_cleanup_task_logs_info_for_skipped_result(monkeypatch, caplog): - monkeypatch.setattr(cleanup_task_module, "_create_agent_backend_client", lambda: object()) - monkeypatch.setattr( - cleanup_task_module, - "cleanup_agent_backend_session", - lambda **kwargs: AgentBackendSessionCleanupResult.skipped("missing_runtime_layer_specs"), - ) - - with caplog.at_level(logging.INFO, logger="tasks.agent_backend_session_cleanup_task"): - cleanup_task_module._run_cleanup_task(_payload_dict()) - - assert "Agent backend session cleanup skipped" in caplog.text - assert "missing_runtime_layer_specs" in caplog.text - - -def test_run_cleanup_task_logs_warning_for_failed_result(monkeypatch, caplog): - monkeypatch.setattr(cleanup_task_module, "_create_agent_backend_client", lambda: object()) - monkeypatch.setattr( - cleanup_task_module, - "cleanup_agent_backend_session", - lambda **kwargs: AgentBackendSessionCleanupResult.failed( - "backend exploded", - cleanup_run_id="cleanup-run-1", - ), - ) - - with caplog.at_level(logging.WARNING, logger="tasks.agent_backend_session_cleanup_task"): - cleanup_task_module._run_cleanup_task(_payload_dict()) - - assert "Agent backend session cleanup failed" in caplog.text - assert "backend exploded" in caplog.text - assert "cleanup-run-1" in caplog.text diff --git a/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py b/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py new file mode 100644 index 00000000000..48fe8ff2def --- /dev/null +++ b/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py @@ -0,0 +1,81 @@ +from typing import Protocol, cast +from unittest.mock import MagicMock + +import pytest + +from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.workspace_service import AgentWorkspaceService +from tasks.collect_agent_resources_task import ( + collect_agent_resources, + enqueue_agent_resource_collection, +) + + +class _TaskWithQueue(Protocol): + queue: str + + +def test_collection_task_uses_retention_queue() -> None: + task = cast(_TaskWithQueue, collect_agent_resources) + assert task.queue == "retention" + + +def test_enqueue_deduplicates_ids_and_skips_empty_input(monkeypatch: pytest.MonkeyPatch) -> None: + delay = MagicMock() + monkeypatch.setattr(collect_agent_resources, "delay", delay) + + enqueue_agent_resource_collection(tenant_id="tenant-1") + enqueue_agent_resource_collection( + tenant_id="tenant-1", + binding_ids=["binding-2", "binding-1", "binding-2"], + workspace_ids=["workspace-1"], + ) + + delay.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1", "binding-2"], + workspace_ids=["workspace-1"], + home_snapshot_ids=[], + ) + + +def test_collection_continues_in_workspace_binding_snapshot_order(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + + def collect_workspace(**_kwargs: object) -> None: + calls.append("workspace") + raise RuntimeError("workspace failed") + + monkeypatch.setattr(AgentWorkspaceService, "collect_retired_workspace", collect_workspace) + monkeypatch.setattr( + AgentWorkspaceService, + "collect_retired_binding", + lambda **_kwargs: calls.append("binding"), + ) + monkeypatch.setattr( + AgentHomeSnapshotService, + "collect_retired_home_snapshot", + lambda **_kwargs: calls.append("home"), + ) + + collect_agent_resources.run( + tenant_id="tenant-1", + workspace_ids=["workspace-1"], + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + ) + + assert calls == ["workspace", "binding", "home"] + + +def test_enqueue_failure_is_best_effort(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + collect_agent_resources, + "delay", + MagicMock(side_effect=RuntimeError("queue unavailable")), + ) + + enqueue_agent_resource_collection( + tenant_id="tenant-1", + binding_ids=["binding-1"], + ) diff --git a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py index 7d2705447db..3f0ce2475ef 100644 --- a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py +++ b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py @@ -1,22 +1,17 @@ import logging -from collections.abc import Generator from datetime import UTC, datetime from unittest.mock import MagicMock, call, patch from uuid import uuid4 import pytest -from agenton.compositor import CompositorSessionSnapshot -from sqlalchemy import delete, select from sqlalchemy.orm import Session -from core.db.session_factory import session_factory from graphon.enums import WorkflowExecutionStatus from libs.archive_storage import ArchiveStorageNotConfiguredError -from models import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus, AppStar +from models import AppStar from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.workflow import WorkflowArchiveLog from tasks.remove_app_and_related_data_task import ( - _cleanup_active_agent_runtime_sessions_for_app, _delete_app_stars, _delete_app_workflow_archive_logs, _delete_archived_workflow_run_files, @@ -26,29 +21,6 @@ from tasks.remove_app_and_related_data_task import ( ) -@pytest.fixture(autouse=True) -def _create_agent_runtime_sessions_table() -> Generator[None, None, None]: - engine = session_factory.get_session_maker().kw["bind"] - AgentRuntimeSession.__table__.create(bind=engine, checkfirst=True) - yield - with session_factory.create_session() as session: - session.execute(delete(AgentRuntimeSession)) - session.commit() - AgentRuntimeSession.__table__.drop(bind=engine, checkfirst=True) - - -def _runtime_session_specs_json() -> str: - return ( - '[{"name":"execution_context","type":"dify.execution_context","deps":{},"metadata":{},' - '"config":{"tenant_id":"tenant-1"}},{"name":"history","type":"pydantic_ai.history","deps":{},' - '"metadata":{},"config":null}]' - ) - - -def _snapshot_json() -> str: - return CompositorSessionSnapshot(layers=[]).model_dump_json() - - class TestDeleteDraftVariablesBatch: def test_delete_draft_variables_batch_invalid_batch_size(self): """Test that invalid batch size raises ValueError.""" @@ -89,18 +61,12 @@ class TestDeleteDraftVariableOffloadData: """Test handling of database operation failures.""" mock_conn = MagicMock() file_ids = ["file-1"] - - # Make execute raise an exception mock_conn.execute.side_effect = Exception("Database error") - # Execute function - should not raise, but log error with caplog.at_level(logging.ERROR): result = _delete_draft_variable_offload_data(mock_conn, file_ids) - # Should return 0 when error occurs assert result == 0 - - # Verify error was logged assert "Error deleting draft variable offload data:" in caplog.text @@ -217,215 +183,3 @@ class TestDeleteArchivedWorkflowRunFiles: storage.list_objects.assert_called_once_with("tenant-1/app_id=app-1/") storage.delete_object.assert_has_calls([call("key-1"), call("key-2")], any_order=False) assert "Deleted 2 archive objects for app app-1" in caplog.text - - -class TestCleanupActiveAgentRuntimeSessionsForApp: - @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") - @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") - def test_enqueues_cleanup_for_active_rows_and_marks_rows_cleaned( - self, - mock_conversation_cleanup, - mock_workflow_cleanup, - ): - with session_factory.create_session() as session: - session.add( - AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-1", - agent_config_snapshot_id="snap-1", - conversation_id="conv-1", - backend_run_id="run-conv", - session_snapshot=_snapshot_json(), - composition_layer_specs=_runtime_session_specs_json(), - status=AgentRuntimeSessionStatus.ACTIVE, - ) - ) - session.add( - AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.WORKFLOW_RUN, - agent_id="agent-2", - workflow_id="wf-1", - workflow_run_id="wf-run-1", - node_id="node-1", - backend_run_id="run-wf", - session_snapshot=_snapshot_json(), - composition_layer_specs=_runtime_session_specs_json(), - status=AgentRuntimeSessionStatus.ACTIVE, - ) - ) - session.add( - AgentRuntimeSession( - tenant_id="tenant-1", - app_id="other-app", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-3", - conversation_id="conv-2", - backend_run_id="run-other", - session_snapshot=_snapshot_json(), - composition_layer_specs=_runtime_session_specs_json(), - status=AgentRuntimeSessionStatus.ACTIVE, - ) - ) - session.commit() - - _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1", batch_size=1) - - assert mock_conversation_cleanup.delay.call_count == 1 - assert mock_workflow_cleanup.delay.call_count == 1 - conversation_payload = mock_conversation_cleanup.delay.call_args.args[0] - workflow_payload = mock_workflow_cleanup.delay.call_args.args[0] - assert conversation_payload["metadata"]["conversation_id"] == "conv-1" - assert workflow_payload["metadata"]["workflow_run_id"] == "wf-run-1" - assert conversation_payload["idempotency_key"].startswith("tenant-1:app-1:conv-1:agent-1:app-delete-cleanup:") - assert workflow_payload["idempotency_key"].startswith( - "tenant-1:app-1:wf-run-1:node-1:agent-2:app-delete-cleanup:" - ) - with session_factory.create_session() as session: - app_rows = session.scalars( - select(AgentRuntimeSession).where( - AgentRuntimeSession.tenant_id == "tenant-1", - AgentRuntimeSession.app_id == "app-1", - ) - ).all() - assert {row.status for row in app_rows} == {AgentRuntimeSessionStatus.CLEANED} - other_row = session.scalar( - select(AgentRuntimeSession).where( - AgentRuntimeSession.tenant_id == "tenant-1", - AgentRuntimeSession.app_id == "other-app", - ) - ) - assert other_row is not None - assert other_row.status == AgentRuntimeSessionStatus.ACTIVE - - @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") - @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") - def test_marks_rows_cleaned_even_when_enqueue_fails( - self, - mock_conversation_cleanup, - mock_workflow_cleanup, - ): - mock_conversation_cleanup.delay.side_effect = RuntimeError("queue down") - with session_factory.create_session() as session: - session.add( - AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-1", - conversation_id="conv-1", - backend_run_id="run-conv", - session_snapshot=_snapshot_json(), - composition_layer_specs=_runtime_session_specs_json(), - status=AgentRuntimeSessionStatus.ACTIVE, - ) - ) - session.add( - AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.WORKFLOW_RUN, - agent_id="agent-2", - workflow_id="wf-1", - workflow_run_id="wf-run-1", - node_id="node-1", - backend_run_id="run-wf", - session_snapshot=_snapshot_json(), - composition_layer_specs=_runtime_session_specs_json(), - status=AgentRuntimeSessionStatus.ACTIVE, - ) - ) - session.commit() - - _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1") - - mock_conversation_cleanup.delay.assert_called_once() - mock_workflow_cleanup.delay.assert_called_once() - with session_factory.create_session() as session: - rows = session.scalars( - select(AgentRuntimeSession).where( - AgentRuntimeSession.tenant_id == "tenant-1", - AgentRuntimeSession.app_id == "app-1", - ) - ).all() - assert rows - assert {row.status for row in rows} == {AgentRuntimeSessionStatus.CLEANED} - - @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") - @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") - def test_uses_row_identity_to_keep_distinct_app_delete_cleanup_jobs_distinct( - self, - mock_conversation_cleanup, - mock_workflow_cleanup, - ): - del mock_workflow_cleanup - with session_factory.create_session() as session: - first = AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-1", - agent_config_snapshot_id="snap-1", - conversation_id="conv-1", - backend_run_id="run-conv-1", - session_snapshot=_snapshot_json(), - composition_layer_specs=_runtime_session_specs_json(), - status=AgentRuntimeSessionStatus.ACTIVE, - ) - second = AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-1", - agent_config_snapshot_id="snap-2", - conversation_id="conv-1", - backend_run_id="run-conv-2", - session_snapshot=_snapshot_json(), - composition_layer_specs=_runtime_session_specs_json(), - status=AgentRuntimeSessionStatus.ACTIVE, - ) - session.add(first) - session.add(second) - session.commit() - - _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1") - - assert mock_conversation_cleanup.delay.call_count == 2 - payloads = [queued_call.args[0] for queued_call in mock_conversation_cleanup.delay.call_args_list] - assert payloads[0]["idempotency_key"] != payloads[1]["idempotency_key"] - assert payloads[0]["idempotency_key"].endswith(first.id) - assert payloads[1]["idempotency_key"].endswith(second.id) - - @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") - @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") - def test_marks_empty_runtime_layer_specs_rows_clean_without_enqueue( - self, - mock_conversation_cleanup, - mock_workflow_cleanup, - ): - del mock_workflow_cleanup - with session_factory.create_session() as session: - row = AgentRuntimeSession( - tenant_id="tenant-1", - app_id="app-1", - owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, - agent_id="agent-1", - conversation_id="conv-1", - backend_run_id="run-no-specs", - session_snapshot=_snapshot_json(), - composition_layer_specs="[]", - status=AgentRuntimeSessionStatus.ACTIVE, - ) - session.add(row) - session.commit() - - _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1") - - mock_conversation_cleanup.delay.assert_not_called() - with session_factory.create_session() as session: - stored_row = session.scalar(select(AgentRuntimeSession).where(AgentRuntimeSession.id == row.id)) - assert stored_row is not None - assert stored_row.status == AgentRuntimeSessionStatus.CLEANED diff --git a/api/tests/unit_tests/tasks/test_resume_agent_app_task.py b/api/tests/unit_tests/tasks/test_resume_agent_app_task.py index 3b1c887e8ca..684ee06edf3 100644 --- a/api/tests/unit_tests/tasks/test_resume_agent_app_task.py +++ b/api/tests/unit_tests/tasks/test_resume_agent_app_task.py @@ -58,6 +58,7 @@ def test_resume_happy_path_account_user_sets_tenant_and_runs(mocker: MockerFixtu gen.return_value.resume_after_form_submission.assert_called_once() kwargs = gen.return_value.resume_after_form_submission.call_args.kwargs assert kwargs["conversation_id"] == "conv-1" + assert kwargs["form_id"] == "form-1" assert kwargs["user"] is account assert kwargs["app_model"] is app assert kwargs["invoke_from"] == InvokeFrom.WEB_APP diff --git a/api/tests/unit_tests/tasks/test_workflow_node_execution_identity.py b/api/tests/unit_tests/tasks/test_workflow_node_execution_identity.py new file mode 100644 index 00000000000..e1e8f8bc0cf --- /dev/null +++ b/api/tests/unit_tests/tasks/test_workflow_node_execution_identity.py @@ -0,0 +1,28 @@ +from datetime import UTC, datetime + +from graphon.entities import WorkflowNodeExecution +from models.workflow import WorkflowNodeExecutionModel +from tasks.workflow_node_execution_tasks import _update_node_execution_from_domain + + +def test_celery_update_preserves_workflow_agent_binding_identity() -> None: + stored = WorkflowNodeExecutionModel( + process_data='{"workflow_agent_binding_id": "workflow-binding-1"}', + ) + incoming = WorkflowNodeExecution( + id="execution-1", + workflow_id="workflow-1", + node_id="node-1", + node_type="agent", + title="Agent", + index=1, + created_at=datetime.now(UTC), + process_data={"retry": True}, + ) + + _update_node_execution_from_domain(stored, incoming) + + assert stored.process_data_dict == { + "retry": True, + "workflow_agent_binding_id": "workflow-binding-1", + } diff --git a/api/uv.lock b/api/uv.lock index 3bcfa208722..7f5eb6db7ca 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1293,6 +1293,7 @@ dependencies = [ [package.metadata] requires-dist = [ + { name = "e2b", marker = "extra == 'server'", specifier = ">=2.34.0,<3.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, { name = "graphon", marker = "extra == 'server'", specifier = "==0.5.2" }, { name = "grpclib", extras = ["protobuf"], marker = "extra == 'grpc'", specifier = ">=0.4.9,<0.5.0" }, diff --git a/dify-agent/.example.env b/dify-agent/.example.env index 4f11bb9747d..925baa21f15 100644 --- a/dify-agent/.example.env +++ b/dify-agent/.example.env @@ -26,22 +26,32 @@ DIFY_AGENT_INNER_API_URL=http://localhost:5001 # Must match API/worker INNER_API_KEY_FOR_PLUGIN, not the generic INNER_API_KEY. DIFY_AGENT_INNER_API_KEY= -# Shell layer -# Shell backend used by the dify.shell layer: "shellctl" (default) or "enterprise". -DIFY_AGENT_SHELL_PROVIDER=shellctl -# shellctl provider: base URL for the shellctl server. Leave empty to disable shell layer use. -DIFY_AGENT_SHELLCTL_ENTRYPOINT= -# Optional bearer token sent to the shellctl server. -DIFY_AGENT_SHELLCTL_AUTH_TOKEN= -# Root directory for per-Agent shell HOME directories. Use a writable path such -# as /tmp/dify-agent-home for local macOS development. -DIFY_AGENT_SHELL_HOME_ROOT=/home -# enterprise provider: sandbox gateway endpoint. Leave empty to disable shell layer use. +# Runtime resources +# Select one coherent Home Snapshot + Execution Binding backend: local, enterprise, or e2b. +DIFY_AGENT_RUNTIME_BACKEND=local +# Local backend: shellctl data-plane URL and optional bearer token. +# Leave the endpoint empty when this server will not provide dify.runtime or resource endpoints. +DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT= +DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN= +# Enterprise resource operations currently fail fast with NotImplementedError. +# These names are retained for the configured Enterprise Gateway boundary. DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_ENDPOINT= -# Optional X-Inner-Api-Key sent to the enterprise sandbox gateway. DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_AUTH_TOKEN= DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_TIMEOUT=30 DIFY_AGENT_ENTERPRISE_SANDBOX_PROXY_TIMEOUT=60 +# E2B backend: API key, prepared shellctl template, and active Binding policy. +DIFY_AGENT_E2B_API_KEY= +DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox +# Maximum continuous active time. Binding resources pause; temporary Home initialization resources are killed. +# This is not a retention TTL for paused resources or immutable snapshots. +DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600 +DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN= +DIFY_AGENT_E2B_SHELLCTL_PORT=5004 +# Maximum whole-file size read through /workspace/files for an Agent Stub ToolFile upload (50 MiB). +# The environment variable keeps its existing SANDBOX name for deployment compatibility. +DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES=52428800 +# JSON array of regex patterns to redact from shell output shown to the agent. +DIFY_AGENT_SHELL_REDACT_PATTERNS= # Agent Stub # Public Agent Stub URL reachable from shellctl-managed remote machines. @@ -62,18 +72,11 @@ DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY # Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))' DIFY_AGENT_API_TOKEN=dify-agent-run-token-for-dev-only -# Shared plugin-daemon HTTP client timeouts and limits. -# Plugin-daemon HTTP connect timeout in seconds. -DIFY_AGENT_PLUGIN_DAEMON_CONNECT_TIMEOUT=10 -# Plugin-daemon HTTP read timeout in seconds. -DIFY_AGENT_PLUGIN_DAEMON_READ_TIMEOUT=600 -# Plugin-daemon HTTP write timeout in seconds. -DIFY_AGENT_PLUGIN_DAEMON_WRITE_TIMEOUT=30 -# Plugin-daemon HTTP connection-pool wait timeout in seconds. -DIFY_AGENT_PLUGIN_DAEMON_POOL_TIMEOUT=10 -# Maximum total plugin-daemon HTTP connections. -DIFY_AGENT_PLUGIN_DAEMON_MAX_CONNECTIONS=100 -# Maximum idle keep-alive plugin-daemon HTTP connections. -DIFY_AGENT_PLUGIN_DAEMON_MAX_KEEPALIVE_CONNECTIONS=20 -# Keep-alive expiry in seconds for idle plugin-daemon HTTP connections. -DIFY_AGENT_PLUGIN_DAEMON_KEEPALIVE_EXPIRY=30 +# Shared outbound HTTP client timeouts and limits. +DIFY_AGENT_OUTBOUND_HTTP_CONNECT_TIMEOUT=10 +DIFY_AGENT_OUTBOUND_HTTP_READ_TIMEOUT=600 +DIFY_AGENT_OUTBOUND_HTTP_WRITE_TIMEOUT=30 +DIFY_AGENT_OUTBOUND_HTTP_POOL_TIMEOUT=10 +DIFY_AGENT_OUTBOUND_HTTP_MAX_CONNECTIONS=100 +DIFY_AGENT_OUTBOUND_HTTP_MAX_KEEPALIVE_CONNECTIONS=20 +DIFY_AGENT_OUTBOUND_HTTP_KEEPALIVE_EXPIRY=30 diff --git a/dify-agent/.gitignore b/dify-agent/.gitignore index f644f2e5451..5f6b9ad4495 100644 --- a/dify-agent/.gitignore +++ b/dify-agent/.gitignore @@ -1 +1,2 @@ dify-aio +site/ diff --git a/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md b/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md new file mode 100644 index 00000000000..7a8b4416f6c --- /dev/null +++ b/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md @@ -0,0 +1,194 @@ +# Runtime resources + +Dify separates persistent product resources from request-time execution: + +- a **Home Snapshot** is immutable Agent-owned Home content; +- a **Workspace** is mutable working data owned by a product scope such as a + conversation, Build Draft, or Workflow run; +- an **Execution Binding** is one materialized Agent participant, including its + private Home and resumable session, attached to a Workspace; +- a **RuntimeLease** is operation-scoped access to the physical Binding. + +`AgentWorkspaceBinding.id` is the participant, materialized Home, and persisted +Agenton session identity. `agent_id` identifies the source Agent. The same +Agent can therefore have multiple active Bindings in one Workspace: +each has an independent Home and session, while all may share Workspace files. + +Home and Workspace are logically independent. A backend may still couple their +physical representation. For example, current E2B maps one Binding and its +Workspace to one E2B resource, while Local can attach multiple materialized +Homes to one shared Workspace. + +## Runtime layer graph + +Agent requests do not expose separate Home, Workspace, or Sandbox layers. Dify +API resolves the Binding selected by the product flow and sends its opaque +backend ref to the `dify.runtime` layer: + +```mermaid +flowchart LR + EC["dify.execution_context
request identity"] + RT["dify.runtime
opaque backend_binding_ref"] + SH["dify.shell
commands, files, and jobs"] + + EC --> SH + RT --> SH +``` + +`DifyRuntimeLayer` calls the selected `ExecutionBindingBackend.acquire()` when +its resource context opens and `release()` when the operation ends. It exposes +the resulting `RuntimeLease` only while that context is active. The layer does +not create, retire, or destroy persistent resources, and it stores no backend +SDK object in an Agenton session snapshot. + +The Shell layer consumes `RuntimeLease.commands`, `RuntimeLease.files`, and +`RuntimeLease.layout`. It tracks only request-local shell job ids and offsets. +Closing a run clears that job state; it does not retire the Binding. + +## State ownership + +Dify API is the lifecycle ledger. It stores three resource records: + +| Record | Meaning | Backend field | +| --- | --- | --- | +| `agent_home_snapshots` | One immutable Home version owned by an Agent. | `snapshot_ref` | +| `agent_workspaces` | One mutable Workspace owned by a product scope. | `backend_workspace_ref` | +| `agent_workspace_bindings` | One materialized participant, private Home, and resumable session attached to a Workspace. | `backend_binding_ref` | + +Backend refs are opaque strings interpreted only by the selected backend +adapter. Dify API stores the latest Agenton session snapshot on the Binding, but +it does not serialize `RuntimeLease`, SDK clients, credentials, or temporary +access tokens. + +Dify Agent does not connect to the Dify product database and has no persistent +resource registry. Its private control-plane endpoints create or destroy +backend resources from requests made by Dify API. Redis run records and event +streams are observability state, not the Home/Workspace/Binding ledger. + +## Creation and execution flow + +Home Snapshot initialization uses `POST /home-snapshots/initialize`. Build Draft +Apply uses `POST /home-snapshots/from-binding`: Dify Agent acquires the exact +source Binding, snapshots its materialized Home through the backend-native +operation, releases the lease, and returns a new opaque snapshot ref. Dify API +then stores a new immutable `agent_home_snapshots` row and records its logical id +on the resulting config version. There is no replay or initialization fallback +when the source Binding is unavailable. + +Before an Agent request, Dify API loads the specific product context. If it has +no associated Binding, Dify API materializes one and saves the Binding id in the +same database transaction. Otherwise it resolves only that Binding and validates +its owner and config/Home generation. Missing, retired, or mismatched Bindings +fail fast; Dify API does not search by Agent, Workspace, candidate count, or +recency, and it does not create a replacement implicitly. + +`POST /execution-bindings` materializes the selected Home Snapshot and returns +opaque Binding and Workspace refs. Every create request represents a new +participant, even when the Agent, Snapshot, config generation, and Workspace +match another Binding. The request composition contains: + +```json +{ + "name": "runtime", + "type": "dify.runtime", + "config": {"backend_binding_ref": "opaque-backend-binding-ref"} +} +``` + +Each Agent request acquires that ref for the duration of the run and releases it +afterward. Local release closes the operation's shellctl connection. E2B release +also pauses the underlying E2B resource with memory preserved. A later request +or Workspace file operation acquires a new lease for the same Binding ref. If a +backend confirms the resource is gone, acquisition fails; it does not create an +empty replacement Workspace. + +## Retirement and collection + +Retirement is a database transition from `ACTIVE` to `RETIRED`. It prevents new +product use without performing network I/O inside the caller's transaction. +Product lifecycle paths commit this transition synchronously. After the +transaction commits, one Celery task asks Dify Agent to destroy the physical +resources. A successful collector deletes the corresponding ledger row; a +failed collector logs the failure and leaves the RETIRED row available for a +future retry or reconciler. + +The unified `collect_agent_resources` task is registered on normal Celery +workers and explicitly uses the existing `retention` queue. Standard workers +already consume that queue, so no dedicated Agent resource worker or new queue +is required. At a Workflow terminal event, the graph layer synchronously retires +and commits the run's Workspaces before enqueueing collection. When a Workflow +change may orphan Workflow-only Agents, the main product transaction commits +first; a fresh session then rechecks effective ownership and retires only Agents +that remain unowned. + +Retiring a final Binding also retires its Workspace. Workspace collection +destroys the physical Workspace through one Binding and then collects remaining +materialized Homes. Home Snapshots are retired when their owning Agent is +retired and are collected only after no draft or config snapshot references +them. Celery performs physical collection only; it does not decide or perform +the initial retirement. Dify Agent itself remains stateless. + +There is currently no age-based TTL, periodic GC, or global orphan reconciler. +Backend destroy operations are idempotent where supported. Dify API does not +perform cross-system compensation after a backend create returns success. Any +later API failure, including Python, flush, or commit failure, may leave a +physical orphan for a future global reconciler. + +Backends still clean up partial resources when a create operation fails before +returning success. For example, E2B kills a Sandbox when its initialization +fails, and Local removes paths created by an incomplete operation. This +backend-local cleanup does not cross the database commit boundary. + +## Workspace file boundary + +Dify API's public file APIs accept a product locator, not a Binding id or +backend ref: a Conversation, a debug Build Draft, or a Workflow Node Execution. +Dify API authorizes that object and resolves its associated active Binding. It +does not select the latest Binding or fall back to another product context. + +The resolved request reaches Dify Agent through its private +`POST /workspace/files/list`, `POST /workspace/files/read`, and +`POST /workspace/files/upload` endpoints. Each operation receives a +`backend_binding_ref`, acquires a fresh RuntimeLease, performs the file action, +and releases the lease. + +`WorkspaceFileService` forwards the request path unchanged to the current +RuntimeLease's file capability. The backend interprets that path in its own +filesystem namespace; the service does not require a Workspace-relative path +or enforce `workspace_dir` containment. `~` and `~/...` can therefore address +the lease's Home. Whether an absolute path or a path containing `..` is +accessible is determined by the backend and its path-isolation policy, such as +Local shellctl and Landlock isolation. Whole-file capture for Agent Stub upload +is bounded by +`DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES`; the environment variable keeps its +existing name even though the route now operates through a Binding. + +`RuntimeLayout.home_dir` and `RuntimeLayout.workspace_dir` are canonical paths +inside the backend execution namespace. They are not host paths, product ids, +or request configuration. Shell commands start in `workspace_dir`, and `HOME` +is forced to `home_dir`. On Local, sibling materialized Homes may exist in the +same shellctl namespace, while path isolation restricts the active lease to its +own Home plus the shared Workspace. + +## Backend support + +| Backend | Home Snapshot operations | Binding operations | Physical relationship | +| --- | --- | --- | --- | +| Local | Supported | Supported, including attaching multiple Bindings to one Workspace | Snapshot directory, per-Binding materialized Home, and Workspace directory are separate. | +| E2B | Supported | Supported without shared-Workspace attachment | Binding and Workspace refs map to the same E2B resource; Home initialization/checkpoint uses E2B snapshots. | +| Enterprise | Not implemented | Not implemented | Configuration is accepted, but every resource operation fails fast with `NotImplementedError`. | + +Local creates a new Home for every Binding id. Destroying one Binding without +the Workspace leaves sibling Homes and the shared Workspace intact. Current E2B +rejects `existing_workspace_ref` with `shared_workspace_unsupported`, because +its Binding and Workspace are one Sandbox. It also rejects binding-only destroy. +Neither path creates a fallback Workspace or switches backends. + +`DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` limits continuous active time for an E2B +resource. Runtime resources pause on timeout; temporary Home initialization +resources are killed. It is not a retention TTL and does not delete paused +resources or immutable snapshots. + +See the [Shell layer](../../user-manual/shell-layer/index.md) for request +composition and the [Operations Guide](../../guide/index.md) for Local and E2B +validation. diff --git a/dify-agent/docs/dify-agent/get-started/index.md b/dify-agent/docs/dify-agent/get-started/index.md index 9c0056d772c..97bb8848383 100644 --- a/dify-agent/docs/dify-agent/get-started/index.md +++ b/dify-agent/docs/dify-agent/get-started/index.md @@ -73,11 +73,37 @@ The minimum settings are: See `.example.env` for the full server settings template. -If you plan to run `dify.shell`, also configure `DIFY_AGENT_SHELLCTL_ENTRYPOINT` -and, when shell jobs need to call back with the `dify-agent` command, set +If you plan to run `dify.shell`, select a coherent Home Snapshot and Execution +Binding backend. A standalone Local server normally uses: + +```env +DIFY_AGENT_RUNTIME_BACKEND=local +DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://127.0.0.1:5004 +DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN= +DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES=52428800 +``` + +E2B requires `DIFY_AGENT_E2B_API_KEY` and defaults to the prepared +`difys-default-team/dify-agent-local-sandbox` template. The E2B active timeout +pauses the physical resource behind a Binding or kills temporary Home +initialization resources; it is not a retention TTL. Enterprise settings are +accepted, but current Home Snapshot and Binding operations fail fast with +`NotImplementedError`. + +A shell-enabled request includes Execution Context, `dify.runtime`, and +`dify.shell`. Dify API creates or resolves the specific persistent Binding for +the request's product context; +`DifyRuntimeLayerConfig.backend_binding_ref` carries only that opaque ref and +opens a new operation-scoped `RuntimeLease` for the run. When shell jobs need to +call back with the `dify-agent` command, also set `DIFY_AGENT_STUB_API_BASE_URL`. The supplied default configs include a development `DIFY_AGENT_SERVER_SECRET_KEY`, but production deployments should -override it with a unique 32-byte base64url value as documented in `.example.env`. +override it with a unique 32-byte base64url value as documented in +`.example.env`. + +See [Runtime resources](../concepts/runtime-resources/index.md) for the layer +graph and state ownership, and the [Operations Guide](../guide/index.md) for +backend-specific configuration and validation commands. ## Start the Dify Agent server diff --git a/dify-agent/docs/dify-agent/guide/index.md b/dify-agent/docs/dify-agent/guide/index.md index d588e5b7960..ca6e739f3a5 100644 --- a/dify-agent/docs/dify-agent/guide/index.md +++ b/dify-agent/docs/dify-agent/guide/index.md @@ -15,7 +15,8 @@ uv run --project dify-agent uvicorn dify_agent.server.app:app --reload By default, the FastAPI lifespan creates: - one Redis-backed run store used by HTTP routes -- one shared plugin-daemon `httpx.AsyncClient` used by local run tasks +- shared plugin-daemon and Dify API inner `httpx.AsyncClient` instances +- one deployment-selected, stateless runtime backend profile when configured - one process-local scheduler that starts background `asyncio` run tasks This means local development needs one uvicorn process plus Redis, and @@ -38,19 +39,30 @@ also reads `.env` and `dify-agent/.env` when present. | `DIFY_AGENT_PLUGIN_DAEMON_API_KEY` | empty | API key sent to the Dify plugin daemon. | | `DIFY_AGENT_INNER_API_URL` | `http://localhost:5001` | Dify API service root used when dify-agent calls `/inner/api/...` endpoints. | | `DIFY_AGENT_INNER_API_KEY` | empty | API key sent to Dify API inner plugin endpoints. Set this to Dify API `INNER_API_KEY_FOR_PLUGIN` (Docker: `PLUGIN_DIFY_INNER_API_KEY`). | -| `DIFY_AGENT_SHELLCTL_ENTRYPOINT` | empty | Base URL for the shellctl server used by `dify.shell`; required when runs include the shell layer. | -| `DIFY_AGENT_SHELLCTL_AUTH_TOKEN` | empty | Optional bearer token sent to the shellctl server. | -| `DIFY_AGENT_SHELL_HOME_ROOT` | `/home` | Root for per-Agent shell HOME directories. Set this to a writable path such as `/tmp/dify-agent-home` for local macOS development. | +| `DIFY_AGENT_RUNTIME_BACKEND` | `local` | Selects one coherent `local`, `enterprise`, or `e2b` Home Snapshot + Execution Binding backend profile. | +| `DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT` | empty | Local shellctl data-plane URL. With the default Local selection, leaving it empty disables `dify.runtime` and resource endpoints. | +| `DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN` | empty | Optional bearer token sent to Local shellctl. | +| `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_ENDPOINT` | empty | Enterprise Gateway endpoint required by configuration. Current Home Snapshot and Binding operations fail fast with `NotImplementedError`. | +| `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_AUTH_TOKEN` | empty | Optional `X-Inner-Api-Key` sent to the Enterprise Gateway. | +| `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_TIMEOUT` | `30` | Enterprise control-plane timeout in seconds. | +| `DIFY_AGENT_ENTERPRISE_SANDBOX_PROXY_TIMEOUT` | `60` | Enterprise shellctl-proxy timeout in seconds. | +| `DIFY_AGENT_E2B_API_KEY` | empty | E2B API key; required for E2B. | +| `DIFY_AGENT_E2B_TEMPLATE` | `difys-default-team/dify-agent-local-sandbox` | Prepared E2B template containing shellctl and the initial Home environment. | +| `DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` | `3600` | Maximum continuous active time, up to 3600 seconds. Binding resources pause on timeout; temporary Home initialization resources are killed. This is not a retention TTL. | +| `DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN` | empty | Optional bearer token expected by shellctl inside the E2B template. | +| `DIFY_AGENT_E2B_SHELLCTL_PORT` | `5004` | shellctl port exposed by the E2B template. | +| `DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES` | `52428800` | Standalone Dify Agent maximum for whole-file Workspace upload capture; 50 MiB by default. Docker Compose derives it from `PLUGIN_MAX_FILE_SIZE`. | +| `DIFY_AGENT_SHELL_REDACT_PATTERNS` | empty | JSON array of additional regex patterns redacted from Shell output. | | `DIFY_AGENT_STUB_API_BASE_URL` | empty | Public Agent Stub API base URL reachable from shellctl-managed remote machines. HTTP may be the service root or `/agent-stub`; gRPC must be `grpc://host:port`. Enables `DIFY_AGENT_STUB_*` env injection for user `shell.run` jobs. | | `DIFY_AGENT_STUB_GRPC_BIND_ADDRESS` | empty | Optional `host:port` bind override used only when `DIFY_AGENT_STUB_API_BASE_URL` uses `grpc://`. | | `DIFY_AGENT_SERVER_SECRET_KEY` | empty | Security-sensitive server-wide root secret used to derive the JWE encryption key for Agent Stub bearer tokens; required when `DIFY_AGENT_STUB_API_BASE_URL` is set. The supplied default config uses a development value; set a unique unpadded base64url 32-byte secret in production. | -| `DIFY_AGENT_PLUGIN_DAEMON_CONNECT_TIMEOUT` | `10` | Plugin-daemon HTTP connect timeout in seconds. | -| `DIFY_AGENT_PLUGIN_DAEMON_READ_TIMEOUT` | `600` | Plugin-daemon HTTP read timeout in seconds. | -| `DIFY_AGENT_PLUGIN_DAEMON_WRITE_TIMEOUT` | `30` | Plugin-daemon HTTP write timeout in seconds. | -| `DIFY_AGENT_PLUGIN_DAEMON_POOL_TIMEOUT` | `10` | Plugin-daemon HTTP connection-pool wait timeout in seconds. | -| `DIFY_AGENT_PLUGIN_DAEMON_MAX_CONNECTIONS` | `100` | Maximum total plugin-daemon HTTP connections. | -| `DIFY_AGENT_PLUGIN_DAEMON_MAX_KEEPALIVE_CONNECTIONS` | `20` | Maximum idle keep-alive plugin-daemon HTTP connections. | -| `DIFY_AGENT_PLUGIN_DAEMON_KEEPALIVE_EXPIRY` | `30` | Keep-alive expiry in seconds for idle plugin-daemon HTTP connections. | +| `DIFY_AGENT_OUTBOUND_HTTP_CONNECT_TIMEOUT` | `10` | Shared outbound HTTP connect timeout in seconds. | +| `DIFY_AGENT_OUTBOUND_HTTP_READ_TIMEOUT` | `600` | Shared outbound HTTP read timeout in seconds. | +| `DIFY_AGENT_OUTBOUND_HTTP_WRITE_TIMEOUT` | `30` | Shared outbound HTTP write timeout in seconds. | +| `DIFY_AGENT_OUTBOUND_HTTP_POOL_TIMEOUT` | `10` | Shared outbound connection-pool wait timeout in seconds. | +| `DIFY_AGENT_OUTBOUND_HTTP_MAX_CONNECTIONS` | `100` | Maximum total shared outbound HTTP connections. | +| `DIFY_AGENT_OUTBOUND_HTTP_MAX_KEEPALIVE_CONNECTIONS` | `20` | Maximum idle shared outbound HTTP connections. | +| `DIFY_AGENT_OUTBOUND_HTTP_KEEPALIVE_EXPIRY` | `30` | Idle keep-alive expiry in seconds. | Example `.env`: @@ -63,9 +75,10 @@ DIFY_AGENT_PLUGIN_DAEMON_URL=http://localhost:5002 DIFY_AGENT_PLUGIN_DAEMON_API_KEY=replace-with-daemon-key DIFY_AGENT_INNER_API_URL=http://localhost:5001 DIFY_AGENT_INNER_API_KEY=replace-with-dify-inner-api-key-for-plugin -DIFY_AGENT_SHELLCTL_ENTRYPOINT=http://127.0.0.1:5004 -DIFY_AGENT_SHELLCTL_AUTH_TOKEN=replace-with-shellctl-token -DIFY_AGENT_SHELL_HOME_ROOT=/tmp/dify-agent-home +DIFY_AGENT_RUNTIME_BACKEND=local +DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://127.0.0.1:5004 +DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN=replace-with-shellctl-token +DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES=52428800 DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. @@ -73,10 +86,144 @@ DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY ``` +`DIFY_AGENT_SHELLCTL_ENTRYPOINT` and `DIFY_AGENT_SHELLCTL_AUTH_TOKEN` remain +accepted only as legacy aliases for the two Local settings. New deployments +must use `DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT` and +`DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN`. There is no compatibility setting for +the removed shell-provider selector or Shell-owned Home root. + +The example above is for a standalone Dify Agent process, where the byte limit +can be set directly. In a Docker deployment, set `PLUGIN_MAX_FILE_SIZE` in +`docker/.env`; Compose maps it to +`DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES` inside `agent_backend`. + +The backend selection is deployment-private. Shell-enabled run requests use an +Execution Context, `dify.runtime`, and `dify.shell` graph. Runtime config carries +only the opaque `backend_binding_ref` resolved by Dify API. See +[Runtime resources](../concepts/runtime-resources/index.md) for the ownership +and lifecycle contract. + Run records and event streams use the same retention. Status writes refresh the record TTL, and event writes refresh both the stream TTL and the corresponding record TTL so active runs that keep producing events remain observable. +## Validate the E2B Compose deployment + +The E2B overlay requires a prepared template that starts shellctl on port 5004. +The default is `difys-default-team/dify-agent-local-sandbox`. It also requires +an E2B API key in `DIFY_AGENT_E2B_API_KEY`; the Compose interpolation accepts +`E2B_API_KEY` and `E2B_API_TOKEN` only as deployment-level fallbacks. + +From the repository root, keep the normal `docker/.env` unchanged and export the +secret for this shell: + +```bash +export DIFY_AGENT_E2B_API_KEY="${DIFY_AGENT_E2B_API_KEY:-${E2B_API_KEY:-${E2B_API_TOKEN:-}}}" +test -n "$DIFY_AGENT_E2B_API_KEY" +export DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox + +docker compose \ + --env-file docker/.env \ + -f docker/docker-compose.yaml \ + -f docker/docker-compose.e2b.yaml \ + up -d --build +``` + +The overlay builds `dify-api:e2b-local` and +`dify-agent-backend:e2b-local` from the current checkout. It disables the +normal Local sandbox service, switches Dify Agent to E2B, and mounts PostgreSQL +on the separate `dify_e2b_postgres_data` Compose volume. That database is empty +when the volume is first created and is isolated from the normal stack; later +starts reuse it until an operator explicitly removes the volume. + +Verify the merged deployment and the branch-built Dify Agent API: + +```bash +docker compose \ + --env-file docker/.env \ + -f docker/docker-compose.yaml \ + -f docker/docker-compose.e2b.yaml \ + ps + +agent_backend_port="$( + docker compose \ + --env-file docker/.env \ + -f docker/docker-compose.yaml \ + -f docker/docker-compose.e2b.yaml \ + port agent_backend 5050 | awk -F: 'NR == 1 { print $NF }' +)" +test -n "$agent_backend_port" +curl --fail --silent --show-error \ + --connect-timeout 2 --max-time 5 \ + --retry 12 --retry-delay 1 --retry-connrefused --retry-max-time 60 \ + "http://127.0.0.1:${agent_backend_port}/openapi.json" \ + >/dev/null + +docker compose \ + --env-file docker/.env \ + -f docker/docker-compose.yaml \ + -f docker/docker-compose.e2b.yaml \ + logs --tail=100 agent_backend api worker +``` + +Stop the validation stack without deleting its isolated database volume: + +```bash +docker compose \ + --env-file docker/.env \ + -f docker/docker-compose.yaml \ + -f docker/docker-compose.e2b.yaml \ + down +``` + +`DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` controls continuous active E2B time. +The physical resource behind a Binding pauses when that timeout fires; a +temporary Home initialization resource is killed. Pausing preserves the current +Workspace. The setting does not delete an aged paused resource or immutable +snapshot, and Dify Agent currently has no resource-age TTL, reconciler, or +eventual cleanup guarantee. Dify API retirement followed by Binding collection +kills the coupled E2B resource. + +## Run runtime-backend integration contracts + +Run the disposable Local contract from the `dify-agent` directory. The script +starts one local-sandbox container on an unused port and removes that exact +container on exit: + +```bash +cd dify-agent +DIFY_AGENT_TEST_LOCAL_SANDBOX_IMAGE=langgenius/dify-agent-local-sandbox:1.16.0 \ + tests/integration/dify_agent/runtime_backend/run_local_integration.sh +``` + +To use an already managed Local shellctl endpoint instead: + +```bash +cd dify-agent +DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT=http://127.0.0.1:5004 \ +DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN=replace-with-shellctl-token \ + pdm run pytest --import-mode=importlib \ + tests/integration/dify_agent/runtime_backend/test_working_environment.py \ + -k local -q -rs +``` + +Run the real E2B contract with an explicit test credential and template: + +```bash +cd dify-agent +DIFY_AGENT_TEST_E2B_API_KEY="$E2B_API_TOKEN" \ +DIFY_AGENT_TEST_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox \ +DIFY_AGENT_TEST_E2B_ACTIVE_TIMEOUT_SECONDS=900 \ + pdm run pytest --import-mode=importlib \ + tests/integration/dify_agent/runtime_backend/test_working_environment.py \ + -k e2b -q -rs +``` + +The Local auth token is optional when shellctl has authentication disabled. +The E2B test timeout value `900` means up to 15 minutes of continuous active +test time; it is not a post-test retention TTL. Both contracts create unique +resources and perform explicit cleanup in `finally` blocks. + ## Scheduling and shutdown semantics `POST /runs` persists a `running` run record and starts an `asyncio` task in the diff --git a/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md b/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md index f3f6a77c4de..168dfaa0c7c 100644 --- a/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md @@ -1,126 +1,124 @@ # Shell layer -The shell layer lets a Dify Agent run expose a `shellctl`-backed workspace to the -model. This page is for Dify Agent clients that build `CreateRunRequest` -payloads. It explains how to add the layer to a run composition and how the -server-side runtime must be wired. +The `dify.shell` layer exposes shellctl-backed commands and files to an Agent. +It does not select a backend or own persistent Home, Workspace, or Binding +resources. It consumes the operation-scoped `RuntimeLease` opened by a sibling +`dify.runtime` layer. -The layer type id is `dify.shell`. Its public config is intentionally empty: +## Public configuration ```python -from dify_agent.layers.shell import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig +from dify_agent.layers.shell import ( + DIFY_SHELL_LAYER_TYPE_ID, + DifyShellEnvVarConfig, + DifyShellLayerConfig, +) from dify_agent.protocol import RunLayerSpec RunLayerSpec( name="shell", type=DIFY_SHELL_LAYER_TYPE_ID, - config=DifyShellLayerConfig(), -) -``` - -Server-only settings, such as the `shellctl` HTTP entrypoint and auth token, are -injected by the Dify Agent runtime provider. They are not part of -`DifyShellLayerConfig` and should not be submitted by clients in the public run -request. - -## Runtime requirements - -When a run includes `dify.shell`, the Dify Agent server must construct its layer -providers with a non-empty shellctl entrypoint: - -```python -from dify_agent.adapters.shell.shellctl import ShellctlProvider -from dify_agent.runtime.compositor_factory import create_default_layer_providers - -layer_providers = create_default_layer_providers( - plugin_daemon_url="http://localhost:5002", - plugin_daemon_api_key="replace-with-plugin-daemon-key", - shell_provider=ShellctlProvider( - entrypoint="http://127.0.0.1:5004", - token="replace-with-shellctl-token", # optional; defaults to empty string + deps={"execution_context": "execution_context", "runtime": "runtime"}, + config=DifyShellLayerConfig( + env=[DifyShellEnvVarConfig(name="REPORT_FORMAT", value="markdown")], + redact_patterns=["private-[A-Za-z0-9]+"], ), ) ``` -In the FastAPI server, these values are read from environment-backed -`ServerSettings` fields: +| Config field | Meaning | +| --- | --- | +| `agent_stub_drive_ref` | Optional Drive ref used by shell-visible Agent Stub commands. | +| `cli_tools` | CLI bootstrap declarations with install commands and scoped environment metadata. | +| `env` | Normal environment variables exported to Shell commands. | +| `secret_refs` | Names of secret environment variables supplied by the backend environment. | +| `redact_patterns` | Request-level regex patterns removed from Shell output shown to the model. | -```env -DIFY_AGENT_SHELLCTL_ENTRYPOINT=http://127.0.0.1:5004 -DIFY_AGENT_SHELLCTL_AUTH_TOKEN=replace-with-shellctl-token -DIFY_AGENT_SHELL_HOME_ROOT=/tmp/dify-agent-home +Endpoints, credentials, Home/Workspace paths, resource refs, timeouts, and +network policy are not Shell config. Backend selection is server-private, and +the opaque Binding ref belongs to `DifyRuntimeLayerConfig`. + +## Runtime requirements + +The server constructs one coherent runtime backend profile. Local and E2B +implement Home Snapshot and Execution Binding operations. Enterprise settings +can be selected, but resource operations currently fail fast with +`NotImplementedError`; there is no compatibility fallback to the retired +Sandbox protocol. + +```python +from dify_agent.runtime.compositor_factory import create_default_layer_providers +from dify_agent.runtime_backend.profile import RuntimeBackendSettings, create_runtime_backend_profile + +runtime_backend_profile = create_runtime_backend_profile( + RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="http://127.0.0.1:5004", + local_sandbox_auth_token="replace-with-shellctl-token", + ) +) + +layer_providers = create_default_layer_providers( + plugin_daemon_url="http://localhost:5002", + plugin_daemon_api_key="replace-with-plugin-daemon-key", + runtime_backend_profile=runtime_backend_profile, +) ``` -`DIFY_AGENT_SHELLCTL_AUTH_TOKEN` defaults to `None`/empty, which keeps the shell -client on the no-token path. Set it only when the shellctl server is started with -bearer authentication. +Equivalent standalone environment settings are: -`DIFY_AGENT_SHELL_HOME_ROOT` defaults to `/home`, matching Linux and container -deployments. For local macOS development, set it to a writable directory such as -`/tmp/dify-agent-home`; the shell layer creates per-Agent home directories under -that root. +```env +DIFY_AGENT_RUNTIME_BACKEND=local +DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://127.0.0.1:5004 +DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN=replace-with-shellctl-token +``` -To let commands inside user-visible shell jobs call back to the Dify Agent server -with `dify-agent ...`, also enable the Agent Stub: +The auth token may be empty when shellctl authentication is disabled. E2B uses +`DIFY_AGENT_E2B_API_KEY`, the prepared template, and its shellctl settings. + +To let shell jobs call the Agent Stub with `dify-agent ...`, configure a public +Agent Stub URL and a unique production secret: ```env DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub -# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. -# Replace this development default in production. -# Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))' -DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY +DIFY_AGENT_SERVER_SECRET_KEY=replace-with-unpadded-base64url-for-32-random-bytes ``` -HTTP `DIFY_AGENT_STUB_API_BASE_URL` may be either the service root or the -explicit `/agent-stub` API root; the server normalizes the service root to -`/agent-stub`. Other HTTP paths are rejected at startup. +HTTP URLs may be either the service root or the explicit `/agent-stub` root. +The server normalizes a service root and rejects unrelated paths. -The supplied Docker and `.example.env` configs use a development -`DIFY_AGENT_SERVER_SECRET_KEY`. Override it in production with unpadded base64url -text for exactly 32 decoded bytes. One way to generate it is: +## Request graph -```bash -python -c 'import secrets; print(secrets.token_urlsafe(32))' +A shell-enabled run contains Execution Context, Runtime, and Shell layers: + +```mermaid +flowchart LR + EC["execution_context"] --> SH["shell"] + RT["runtime
backend_binding_ref"] --> SH ``` -## Client request shape +`DifyRuntimeLayer` acquires the Binding when the run's resource context opens +and releases it when that operation exits. `DifyShellLayer` uses the active +lease's commands, files, Home path, and Workspace path. It performs only +best-effort cleanup of shell jobs; the persistent Binding lifecycle remains in +Dify API. -A client adds the shell layer as an ordinary composition layer. Basic shell jobs -do not need dependencies. To inject `DIFY_AGENT_STUB_API_BASE_URL`, -`DIFY_AGENT_STUB_AUTH_JWE`, and `DIFY_AGENT_STUB_DRIVE_BASE` into user-visible -`shell.run` jobs, declare the execution-context layer as the shell layer's -`execution_context` dependency. When the run also includes `dify.drive`, declare -it as the shell layer's `drive` dependency; the injected drive base is then -computed from the fixed Agent Stub drive mount and the drive reference, for -example `/mnt/drive/agent-123`. Without a drive dependency, the CLI keeps the -historical `/mnt/drive` fallback. A typical run still also includes: +## Example request -- a prompt layer that supplies the task; -- an execution-context layer carrying tenant/user context; -- an LLM layer named `llm`. - -When clients want the shell workspace and shellctl job records to be removed at -the end of the run, set `on_exit.default` to `delete`. - -## Example: CSV analysis run - -The following example mirrors a verified run with a real Gemini model and a -temporary shellctl server. The client gives the model a small CSV-shaped dataset -and asks for computed metrics without prescribing the exact shell commands. - -### Request +The Binding must already have been resolved or created by Dify API. Its backend +ref is opaque to the request builder: ```python {test="skip" lint="skip"} -from agenton.layers import ExitIntent from agenton_collections.layers.plain import PromptLayerConfig from dify_agent.layers.dify_plugin.configs import DifyPluginLLMLayerConfig from dify_agent.layers.execution_context import ( DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, DifyExecutionContextLayerConfig, ) +from dify_agent.layers.runtime import DIFY_RUNTIME_LAYER_TYPE_ID, DifyRuntimeLayerConfig from dify_agent.layers.shell import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig from dify_agent.protocol import DIFY_AGENT_MODEL_LAYER_ID -from dify_agent.protocol.schemas import CreateRunRequest, LayerExitSignals, RunComposition, RunLayerSpec +from dify_agent.protocol.schemas import CreateRunRequest, RunComposition, RunLayerSpec request = CreateRunRequest( @@ -130,35 +128,33 @@ request = CreateRunRequest( name="prompt", type="plain.prompt", config=PromptLayerConfig( - prefix="You are a practical data analyst. Give a concise final answer.", - user="""Analyze this small sales dataset with pandas. Use any local computation you think is useful. - -region,product,units,unit_price -north,widget,12,3.50 -north,gadget,5,9.00 -south,widget,7,3.50 -south,gadget,9,9.00 -west,widget,4,3.50 -west,gadget,11,9.00 - -Report the total revenue, the region with the most revenue, total units by -product, and a SHA-256 hash of the CSV content.""", + prefix="Use the workspace when local computation helps.", + user="Create report.txt containing the current UTC timestamp, then summarize it.", ), ), - RunLayerSpec( - name="shell", - type=DIFY_SHELL_LAYER_TYPE_ID, - deps={"execution_context": "execution_context"}, - config=DifyShellLayerConfig(), - ), RunLayerSpec( name="execution_context", type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, config=DifyExecutionContextLayerConfig( tenant_id="92cca973-2d6f-45e0-906e-0b7eda5f2ccf", - invoke_from="workflow_run", + agent_id="8d542564-159d-4168-985c-dde8d8ff6092", + agent_config_version_id="931a4cee-4434-4c1c-8fbd-0a3c7591095d", + agent_config_version_kind="snapshot", + agent_mode="workflow_run", + invoke_from="debugger", ), ), + RunLayerSpec( + name="runtime", + type=DIFY_RUNTIME_LAYER_TYPE_ID, + config=DifyRuntimeLayerConfig(backend_binding_ref="opaque-backend-binding-ref"), + ), + RunLayerSpec( + name="shell", + type=DIFY_SHELL_LAYER_TYPE_ID, + deps={"execution_context": "execution_context", "runtime": "runtime"}, + config=DifyShellLayerConfig(), + ), RunLayerSpec( name=DIFY_AGENT_MODEL_LAYER_ID, type="dify.plugin.llm", @@ -166,88 +162,55 @@ product, and a SHA-256 hash of the CSV content.""", config=DifyPluginLLMLayerConfig( plugin_id="langgenius/gemini", model_provider="google", - model="gemini-3.5-flash", + model="gemini-2.5-flash", credentials={"google_api_key": ""}, ), ), ] - ), - on_exit=LayerExitSignals(default=ExitIntent.DELETE), + ) ) ``` -The same request serialized as JSON has these important layer entries: +The resource part serializes as: ```json { - "composition": { - "schema_version": 1, - "layers": [ - {"name": "prompt", "type": "plain.prompt"}, - { - "name": "shell", - "type": "dify.shell", - "deps": {"execution_context": "execution_context"}, - "config": {} - }, - {"name": "execution_context", "type": "dify.execution_context"}, - { - "name": "llm", - "type": "dify.plugin.llm", - "deps": {"execution_context": "execution_context"} - } - ] - }, - "on_exit": {"default": "delete", "layers": {}} + "layers": [ + { + "name": "runtime", + "type": "dify.runtime", + "config": {"backend_binding_ref": "opaque-backend-binding-ref"} + }, + { + "name": "shell", + "type": "dify.shell", + "deps": {"execution_context": "execution_context", "runtime": "runtime"}, + "config": {} + } + ] } ``` -### Final answer +## Paths and persistence -The terminal `run_succeeded` output was: +`RuntimeLease.layout.home_dir` and `workspace_dir` are absolute paths inside the +backend execution namespace. They are not host filesystem paths and are not +sent in the run request. Shell commands start in `workspace_dir`, while `HOME` +is forced to `home_dir`; `~` therefore resolves to the current Binding's +materialized Home. -````markdown -Here is the analysis of the sales dataset: +Workspace files persist with the Workspace until Dify API retires and collects +it. Releasing a RuntimeLease ends only the current operation. Dify API can later +browse the current Workspace through Dify Agent's private +`/workspace/files/list`, `/workspace/files/read`, and +`/workspace/files/upload` routes, each of which acquires a fresh lease. -* **Total Revenue:** **$305.50** -* **Top Region:** **west** with **$113.00** -* **Total Units by Product:** gadget: 25 units, widget: 23 units -* **SHA-256 Hash:** `e86521a0d759037a09b059cb3cb2419f0a3f06e674db8151ccf2f93811dac0b8` -```` +On Local, multiple Bindings may share a Workspace while each receives a +separate materialized Home. Those directories may be siblings in one shellctl +namespace; path isolation restricts a lease to its Home and Workspace. On E2B, +one physical E2B resource currently represents both Binding and Workspace, so +shared Workspace attachment is unsupported. -## Running the local sandbox in Docker - -Build the local sandbox image from the Dify Agent package root: - -```bash -docker build -f docker/local-sandbox/Dockerfile -t dify-agent-local-sandbox:local . -``` - -Run it with a bearer token and publish the API on localhost: - -```bash -docker run --rm --name dify-agent-local-sandbox \ - -e SHELLCTL_AUTH_TOKEN=replace-with-a-token \ - -p 127.0.0.1:5004:5004 \ - dify-agent-local-sandbox:local -``` - -The image starts `shellctl serve --listen 0.0.0.0:5004` as the non-root -`dify` user. It also sets the fallback `DIFY_AGENT_STUB_DRIVE_BASE=/mnt/drive` -and pre-creates that directory with write access for the same user. - -## Docker image contents - -The provided `docker/local-sandbox/Dockerfile` installs: - -- `tmux`, required by `shellctl` to manage shell jobs; -- common shell workspace tools: `git`, `openssh-client`, `jq`, `ripgrep`, - `unzip`, `zip`, `file`, `procps`, and `less`; -- `dify-agent[grpc,shellctl-server]` as a standalone uv tool, which provides - both the Agent Stub client CLI and the built-in `shellctl` CLI/server; -- `uv`, so uv shebang scripts with PEP 723 metadata can run inside the shell - workspace and Python CLI tools can be installed with isolated tool - environments; -- `node==22.22.1` and `pnpm==11.9.0`, so JavaScript and TypeScript tooling can - run inside the shell workspace without per-job installation; -- a non-root default user named `dify`. +See [Runtime resources](../../concepts/runtime-resources/index.md) for the +ledger and lifecycle contract. The [Operations Guide](../../guide/index.md) +covers Local and E2B validation. diff --git a/dify-agent/mkdocs.yml b/dify-agent/mkdocs.yml index 3993b3618e5..34ba0e2c92d 100644 --- a/dify-agent/mkdocs.yml +++ b/dify-agent/mkdocs.yml @@ -16,6 +16,7 @@ nav: - Overview: dify-agent/index.md - Concepts: - Agent Run Lifecycle: dify-agent/concepts/run-lifecycle/index.md + - Runtime Resources: dify-agent/concepts/runtime-resources/index.md - User Manual: - Get Started: dify-agent/get-started/index.md - Prompt Layer: dify-agent/user-manual/prompt-layer/index.md diff --git a/dify-agent/pyproject.toml b/dify-agent/pyproject.toml index 50bc1a93165..36757323c6b 100644 --- a/dify-agent/pyproject.toml +++ b/dify-agent/pyproject.toml @@ -19,6 +19,7 @@ dify-agent-stub-server = "dify_agent.agent_stub.server.cli:main" [project.optional-dependencies] grpc = ["grpclib[protobuf]>=0.4.9,<0.5.0", "protobuf>=6.33.5,<7.0.0"] server = [ + "e2b>=2.34.0,<3.0.0", "fastapi==0.136.0", "graphon==0.5.2", "jsonschema>=4.23.0,<5.0.0", @@ -47,6 +48,7 @@ extraPaths = ["src", "examples/agenton", "examples/dify_agent"] [tool.pytest.ini_options] testpaths = ["tests"] python_files = ["test_*.py", "*_test.py"] +markers = ["integration: requires a real external service or exercises multiple concrete adapters"] [tool.ruff] line-length = 120 diff --git a/dify-agent/src/dify_agent/adapters/shell/__init__.py b/dify-agent/src/dify_agent/adapters/shell/__init__.py index faeb0cfb808..251c0462ae0 100644 --- a/dify-agent/src/dify_agent/adapters/shell/__init__.py +++ b/dify-agent/src/dify_agent/adapters/shell/__init__.py @@ -12,20 +12,10 @@ from dify_agent.adapters.shell.protocols import ( ShellFileTransferProtocol, ShellPromptObservation, ShellProviderError, - ShellProviderProtocol, - ShellResourceProtocol, ) def __getattr__(name: str) -> object: - if name == "ShellAdapterSettings": - from dify_agent.adapters.shell.config import ShellAdapterSettings - - return ShellAdapterSettings - if name == "create_shell_provider": - from dify_agent.adapters.shell.factory import create_shell_provider - - return create_shell_provider if name == "shellctl": from importlib import import_module @@ -35,14 +25,10 @@ def __getattr__(name: str) -> object: __all__ = [ "CompleteShellCommandResult", - "ShellAdapterSettings", "ShellCommandProtocol", "ShellCommandResult", "ShellCommandStatus", "ShellFileTransferProtocol", "ShellPromptObservation", "ShellProviderError", - "ShellProviderProtocol", - "ShellResourceProtocol", - "create_shell_provider", ] diff --git a/dify-agent/src/dify_agent/adapters/shell/config.py b/dify-agent/src/dify_agent/adapters/shell/config.py deleted file mode 100644 index 3d6140a7fbc..00000000000 --- a/dify-agent/src/dify_agent/adapters/shell/config.py +++ /dev/null @@ -1,61 +0,0 @@ -from typing import ClassVar, Literal, Self -from urllib.parse import urlparse - -from pydantic import Field, model_validator -from pydantic_settings import BaseSettings, SettingsConfigDict - -DEFAULT_SHELL_PROVIDER = "shellctl" - - -class ShellAdapterSettings(BaseSettings): - """Env-backed settings used to construct a shell provider. - - ``shellctl_auth_token`` defaults to ``None``; the factory forwards an empty - string to the shellctl client so it does not fall back to ambient process - credentials. Deployments that enable shellctl bearer auth must set - ``DIFY_AGENT_SHELLCTL_AUTH_TOKEN`` explicitly. - """ - - shell_provider: Literal["shellctl", "enterprise"] = "shellctl" - shellctl_entrypoint: str | None = None - shellctl_auth_token: str | None = None - - enterprise_sandbox_gateway_endpoint: str | None = None - enterprise_sandbox_gateway_auth_token: str | None = None - enterprise_sandbox_gateway_timeout: float = Field(default=30.0, gt=0) - enterprise_sandbox_proxy_timeout: float = Field(default=60.0, gt=0) - - model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict( - env_prefix="DIFY_AGENT_", - env_file=(".env", "dify-agent/.env"), - extra="ignore", - populate_by_name=True, - ) - - @model_validator(mode="after") - def validate_provider_fields(self) -> Self: - match self.shell_provider: - case "shellctl": - if not self.shellctl_entrypoint or not self.shellctl_entrypoint.strip(): - raise ValueError("shellctl_entrypoint is required when shell_provider is 'shellctl'.") - _validate_url(self.shellctl_entrypoint, field_name="shellctl_entrypoint") - case "enterprise": - if not self.enterprise_sandbox_gateway_endpoint or not self.enterprise_sandbox_gateway_endpoint.strip(): - raise ValueError( - "enterprise_sandbox_gateway_endpoint is required when shell_provider is 'enterprise'." - ) - _validate_url( - self.enterprise_sandbox_gateway_endpoint, field_name="enterprise_sandbox_gateway_endpoint" - ) - return self - - -def _validate_url(value: str, *, field_name: str) -> None: - parsed = urlparse(value.strip()) - if parsed.scheme not in ("http", "https") or not parsed.netloc: - raise ValueError(f"{field_name} must be a valid http(s) URL, got: {value!r}") - - -__all__ = [ - "ShellAdapterSettings", -] diff --git a/dify-agent/src/dify_agent/adapters/shell/enterprise/__init__.py b/dify-agent/src/dify_agent/adapters/shell/enterprise/__init__.py deleted file mode 100644 index 71e609c8c54..00000000000 --- a/dify-agent/src/dify_agent/adapters/shell/enterprise/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from dify_agent.adapters.shell.enterprise.enterprise import EnterpriseShellProvider - -__all__ = ["EnterpriseShellProvider"] diff --git a/dify-agent/src/dify_agent/adapters/shell/enterprise/enterprise.py b/dify-agent/src/dify_agent/adapters/shell/enterprise/enterprise.py deleted file mode 100644 index c56b4c32058..00000000000 --- a/dify-agent/src/dify_agent/adapters/shell/enterprise/enterprise.py +++ /dev/null @@ -1,235 +0,0 @@ -from __future__ import annotations - -import logging -from dataclasses import dataclass, field -from typing import TypedDict - -import httpx2 as httpx - -from dify_agent.adapters.shell.protocols import ( - SandboxExpiredError, - ShellCommandProtocol, - ShellFileTransferProtocol, - ShellProviderError, - ShellProviderProtocol, - ShellResourceProtocol, -) -from dify_agent.adapters.shell.shellctl import ( - ShellctlClientProtocol, - ShellctlCommands, - ShellctlFileTransfer, -) - -logger = logging.getLogger(__name__) - - -class _CreateSandboxReply(TypedDict): - sandboxId: str - status: str - - -# --------------------------------------------------------------------------- -# Gateway client (control plane) -# --------------------------------------------------------------------------- - - -@dataclass(slots=True) -class EnterpriseGatewayClient: - """HTTP client for the enterprise sandbox gateway control-plane API.""" - - endpoint: str - auth_token: str - timeout: float = 30.0 - _client: httpx.AsyncClient = field(init=False) - - def __post_init__(self) -> None: - headers: dict[str, str] = {} - if self.auth_token: - headers["X-Inner-Api-Key"] = self.auth_token - self._client = httpx.AsyncClient( - base_url=self.endpoint.rstrip("/"), - headers=headers, - timeout=httpx.Timeout(self.timeout), - ) - - async def create_sandbox(self, *, tenant_id: str | None = None, template: str | None = None) -> _CreateSandboxReply: - body: dict[str, str] = {} - if tenant_id: - body["tenantId"] = tenant_id - if template: - body["template"] = template - response = await self._request("POST", "/v1/sandboxes", json=body) - return response.json() - - async def delete_sandbox(self, sandbox_id: str) -> None: - await self._request("DELETE", f"/v1/sandboxes/{sandbox_id}") - - async def close(self) -> None: - await self._client.aclose() - - async def _request(self, method: str, path: str, **kwargs: object) -> httpx.Response: - try: - response = await self._client.request(method, path, **kwargs) - response.raise_for_status() - return response - except httpx.TimeoutException as exc: - raise ShellProviderError(f"Gateway request timed out: {method} {path}", code="timeout") from exc - except httpx.HTTPStatusError as exc: - raise ShellProviderError( - f"Gateway returned {exc.response.status_code}: {exc.response.text}", - code="gateway_error", - ) from exc - except httpx.RequestError as exc: - raise ShellProviderError(f"Gateway request failed: {exc}", code="request_error") from exc - - -# --------------------------------------------------------------------------- -# Resource & Provider -# --------------------------------------------------------------------------- - - -@dataclass(slots=True) -class EnterpriseResource(ShellResourceProtocol): - """A live enterprise sandbox session. - - Holds the gateway client, sandbox ID, and the shellctl client for data-plane - operations. ``suspend()`` closes the shellctl and gateway clients without - deleting the sandbox, so the same sandbox can be re-attached later. - ``delete()`` additionally calls the gateway to destroy the sandbox pod. - """ - - _sandbox_id: str - gateway: EnterpriseGatewayClient - shellctl_client: ShellctlClientProtocol - commands: ShellCommandProtocol - files: ShellFileTransferProtocol - - @property - def sandbox_id(self) -> str | None: - return self._sandbox_id - - async def suspend(self) -> None: - try: - await self.shellctl_client.close() - except Exception as exc: - logger.warning("Failed to close shellctl client for sandbox %s: %s", self._sandbox_id, exc) - try: - await self.gateway.close() - except Exception as exc: - logger.warning("Failed to close gateway client: %s", exc) - - async def delete(self) -> None: - try: - await self.shellctl_client.close() - except Exception as exc: - logger.warning("Failed to close shellctl client for sandbox %s: %s", self._sandbox_id, exc) - try: - logger.info("Deleting enterprise sandbox via gateway: id=%s", self._sandbox_id) - await self.gateway.delete_sandbox(self._sandbox_id) - logger.info("Enterprise sandbox deleted: id=%s", self._sandbox_id) - except ShellProviderError as exc: - logger.warning("Failed to delete sandbox %s via gateway: %s", self._sandbox_id, exc) - try: - await self.gateway.close() - except Exception as exc: - logger.warning("Failed to close gateway client: %s", exc) - - -@dataclass(slots=True) -class EnterpriseShellProvider(ShellProviderProtocol): - """Provisions enterprise sandboxes via the gateway, connects via shellctl. - - Lifecycle: - 1. ``create()`` calls the gateway to provision a new sandbox pod. - 2. ``attach(sandbox_id)`` builds a shellctl client pointed at the - gateway's ``/proxy/{sandboxId}`` route for an *existing* sandbox, - without provisioning a new one. - 3. Both return an ``EnterpriseResource`` with shellctl-backed - command/file adapters. - 4. ``suspend()`` on the resource closes the shellctl and gateway - clients but leaves the sandbox pod alive. - 5. ``delete()`` on the resource closes clients and deletes the - sandbox via the gateway. - """ - - gateway_endpoint: str - auth_token: str - tenant_id: str | None = None - template: str | None = None - gateway_timeout: float = 30.0 - proxy_timeout: float = 60.0 - - async def create(self) -> EnterpriseResource: - gateway = EnterpriseGatewayClient( - endpoint=self.gateway_endpoint, - auth_token=self.auth_token, - timeout=self.gateway_timeout, - ) - try: - logger.info("Creating enterprise sandbox via gateway %s", self.gateway_endpoint) - reply = await gateway.create_sandbox(tenant_id=self.tenant_id, template=self.template) - sandbox_id = reply["sandboxId"] - logger.info("Enterprise sandbox created: id=%s status=%s", sandbox_id, reply.get("status")) - except BaseException: - await gateway.close() - raise - return self._build_resource(sandbox_id=sandbox_id, gateway=gateway) - - async def attach(self, sandbox_id: str) -> EnterpriseResource: - gateway = EnterpriseGatewayClient( - endpoint=self.gateway_endpoint, - auth_token=self.auth_token, - timeout=self.gateway_timeout, - ) - logger.info("Attaching to existing enterprise sandbox: id=%s", sandbox_id) - resource = self._build_resource(sandbox_id=sandbox_id, gateway=gateway) - # Verify the sandbox is alive by running a trivial command through the proxy. - # If the sandbox has expired, the gateway returns a 404 with "sandbox_expired"; - # if the pod is already gone the gateway returns a 404 with reason "NOT_FOUND". - # Both surface as a ShellProviderError from the shellctl client and mean the - # sandbox must be re-created. - try: - await resource.commands.run("true", timeout=5.0) - except ShellProviderError as exc: - await resource.suspend() - message = str(exc).lower() - if "expired" in message or "not_found" in message: - raise SandboxExpiredError(sandbox_id, cause=exc) from exc - raise - return resource - - def _build_resource(self, *, sandbox_id: str, gateway: EnterpriseGatewayClient) -> EnterpriseResource: - proxy_base_url = f"{self.gateway_endpoint.rstrip('/')}/proxy/" - headers: dict[str, str] = {"X-Sandbox-Id": sandbox_id} - if self.auth_token: - headers["X-Inner-Api-Key"] = self.auth_token - proxy_http_client = httpx.AsyncClient( - base_url=proxy_base_url, - headers=headers, - follow_redirects=True, - timeout=httpx.Timeout(self.proxy_timeout), - transport=httpx.AsyncHTTPTransport(retries=3), - ) - - from shellctl.client import ShellctlClient - - client: ShellctlClientProtocol = ShellctlClient( - proxy_base_url, - token=self.auth_token, - client=proxy_http_client, - ) - - return EnterpriseResource( - _sandbox_id=sandbox_id, - gateway=gateway, - shellctl_client=client, - commands=ShellctlCommands(client=client), - files=ShellctlFileTransfer(client=client), - ) - - -__all__ = [ - "EnterpriseGatewayClient", - "EnterpriseResource", - "EnterpriseShellProvider", -] diff --git a/dify-agent/src/dify_agent/adapters/shell/factory.py b/dify-agent/src/dify_agent/adapters/shell/factory.py deleted file mode 100644 index d5177d96118..00000000000 --- a/dify-agent/src/dify_agent/adapters/shell/factory.py +++ /dev/null @@ -1,31 +0,0 @@ -from dify_agent.adapters.shell.config import ShellAdapterSettings -from dify_agent.adapters.shell.enterprise import EnterpriseShellProvider -from dify_agent.adapters.shell.protocols import ShellProviderProtocol -from dify_agent.adapters.shell.shellctl import ShellctlProvider - - -def create_shell_provider(settings: ShellAdapterSettings | None = None) -> ShellProviderProtocol: - """Return the shell provider selected by ``DIFY_AGENT_SHELL_PROVIDER``.""" - resolved = settings or ShellAdapterSettings() - provider = resolved.shell_provider - match provider: - case "shellctl": - entrypoint = (resolved.shellctl_entrypoint or "").strip() - if not entrypoint: - raise ValueError("DIFY_AGENT_SHELLCTL_ENTRYPOINT is required for the 'shellctl' shell provider.") - return ShellctlProvider( - entrypoint=entrypoint, - token=resolved.shellctl_auth_token or "", - ) - case "enterprise": - return EnterpriseShellProvider( - gateway_endpoint=(resolved.enterprise_sandbox_gateway_endpoint or "").strip(), - auth_token=resolved.enterprise_sandbox_gateway_auth_token or "", - gateway_timeout=resolved.enterprise_sandbox_gateway_timeout, - proxy_timeout=resolved.enterprise_sandbox_proxy_timeout, - ) - case _: - raise ValueError(f"Unknown shell provider: {resolved.shell_provider!r}.") - - -__all__ = ["create_shell_provider"] diff --git a/dify-agent/src/dify_agent/adapters/shell/protocols.py b/dify-agent/src/dify_agent/adapters/shell/protocols.py index 6c9988f0fbd..f5f4e827cee 100644 --- a/dify-agent/src/dify_agent/adapters/shell/protocols.py +++ b/dify-agent/src/dify_agent/adapters/shell/protocols.py @@ -47,19 +47,12 @@ class ShellPromptObservation: class ShellProviderError(RuntimeError): code: str | None + status_code: int | None - def __init__(self, message: str, *, code: str | None = None) -> None: + def __init__(self, message: str, *, code: str | None = None, status_code: int | None = None) -> None: super().__init__(message) self.code = code - - -class SandboxExpiredError(ShellProviderError): - """Raised by a shell provider when ``attach()`` targets a sandbox that no longer exists.""" - - def __init__(self, sandbox_id: str, *, cause: ShellProviderError) -> None: - super().__init__(str(cause), code=cause.code) - self.sandbox_id = sandbox_id - self.__cause__ = cause + self.status_code = status_code class ShellCommandProtocol(Protocol): @@ -118,44 +111,3 @@ class ShellFileTransferProtocol(Protocol): async def upload(self, *, content: bytes, remote_path: str, cwd: str | None = None) -> None: ... async def download(self, *, remote_path: str, cwd: str | None = None) -> bytes: ... - - -class ShellResourceProtocol(Protocol): - @property - def commands(self) -> ShellCommandProtocol: ... - - @property - def files(self) -> ShellFileTransferProtocol: ... - - @property - def sandbox_id(self) -> str | None: ... - - async def suspend(self) -> None: - """Detach from the sandbox without destroying it. - - Called when the resource scope exits with suspend intent. The sandbox - remains alive and can be re-attached later via ``attach(sandbox_id)``. - """ - ... - - async def delete(self) -> None: - """Destroy the sandbox and release all resources. - - Called when the resource scope exits with delete intent. The sandbox is - permanently removed and cannot be re-attached. - """ - ... - - -class ShellProviderProtocol(Protocol): - async def create(self) -> ShellResourceProtocol: - """Provision a new sandbox and return a live resource.""" - ... - - async def attach(self, sandbox_id: str) -> ShellResourceProtocol: - """Connect to an existing sandbox without provisioning a new one. - - The returned resource carries the same ``sandbox_id`` so the caller can - persist it across runs and re-attach as needed. - """ - ... diff --git a/dify-agent/src/dify_agent/adapters/shell/shellctl.py b/dify-agent/src/dify_agent/adapters/shell/shellctl.py index f285483bda5..24cc9a576f8 100644 --- a/dify-agent/src/dify_agent/adapters/shell/shellctl.py +++ b/dify-agent/src/dify_agent/adapters/shell/shellctl.py @@ -1,17 +1,21 @@ -"""Shellctl-backed shell provider adapter for dify-agent. +"""Shellctl command and file data-plane adapters for RuntimeLease objects. The built-in shellctl SDK owns the HTTP timeout policy for long-polling -shellctl requests. This adapter stays narrowly focused on translating SDK and -transport failures into ``ShellProviderError`` so the shell layer can return -tool observations instead of aborting the agent loop. +shellctl requests. This adapter translates SDK and transport failures into +``ShellProviderError``. Browse paths are interpreted directly inside the +backend-provided filesystem namespace. Whole-file reads return bytes to the +control plane; they never create another runtime-visible upload path. """ from __future__ import annotations import base64 import binascii +import json import logging +import posixpath import re +import shlex import time from collections.abc import Awaitable from collections.abc import Callable @@ -19,6 +23,7 @@ from dataclasses import dataclass from typing import Protocol, TypeVar, cast import httpx2 as httpx +from shellctl.client import ShellctlClientError from dify_agent.adapters.shell.protocols import ( ShellCommandProtocol, @@ -26,8 +31,17 @@ from dify_agent.adapters.shell.protocols import ( ShellCommandStatus, ShellFileTransferProtocol, ShellProviderError, - ShellProviderProtocol, - ShellResourceProtocol, +) +from dify_agent.runtime_backend.errors import ( + WorkspaceFileTooLargeError, + WorkspacePathError, + WorkspaceUnavailableError, +) +from dify_agent.runtime_backend.protocols import ( + WorkspaceFileEntry, + WorkspaceFileContent, + WorkspaceListResult, + WorkspaceReadResult, ) logger = logging.getLogger(__name__) @@ -42,6 +56,138 @@ _SHELLCTL_OUTPUT_LIMIT_BYTES = 16 * 1024 _TRANSFER_BEGIN = "<<>>" _TRANSFER_END = "<<>>" _DOWNLOAD_MISSING_EXIT_CODE = 66 +_WORKSPACE_PAYLOAD_BEGIN = "<<>>" +_WORKSPACE_PAYLOAD_END = "<<>>" + +_LIST_WORKSPACE_SCRIPT = r""" +import base64 +import json +import os +import stat +import sys + +path = sys.argv[1] +limit = int(sys.argv[2]) +directory_fd = None +try: + directory_fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) + # DIFY_WORKSPACE_CHECKPOINT: directory_opened + names = sorted(os.listdir(directory_fd)) + entries = [] + for name in names[:limit]: + child_stat = os.stat(name, dir_fd=directory_fd, follow_symlinks=False) + mode = child_stat.st_mode + entry_type = ( + "symlink" if stat.S_ISLNK(mode) else + "dir" if stat.S_ISDIR(mode) else + "file" if stat.S_ISREG(mode) else + "other" + ) + entries.append({ + "name": name, + "type": entry_type, + "size": int(child_stat.st_size), + "mtime": int(child_stat.st_mtime), + }) +finally: + if directory_fd is not None: + os.close(directory_fd) + +payload = {"path": path, "entries": entries, "truncated": len(names) > limit} +blob = base64.b64encode(json.dumps(payload).encode()).decode() +print("<<>>" + blob + "<<>>") +""" + +_READ_WORKSPACE_SCRIPT = r""" +import base64 +import json +import os +import stat +import sys + +path = sys.argv[1] +max_bytes = int(sys.argv[2]) +file_fd = None +try: + file_fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) + # DIFY_WORKSPACE_CHECKPOINT: file_opened + file_stat = os.fstat(file_fd) + if not stat.S_ISREG(file_stat.st_mode): + raise FileNotFoundError(path) + size = int(file_stat.st_size) + data = os.read(file_fd, max_bytes + 1) +finally: + if file_fd is not None: + os.close(file_fd) + +truncated = len(data) > max_bytes +data = data[:max_bytes] +try: + text = data.decode("utf-8") + binary = False +except UnicodeDecodeError: + text = None + binary = True +payload = { + "path": path, + "size": size, + "truncated": truncated, + "binary": binary, + "text": text, +} +blob = base64.b64encode(json.dumps(payload).encode()).decode() +print("<<>>" + blob + "<<>>") +""" + +_READ_WORKSPACE_BYTES_SCRIPT = r""" +import base64 +import json +import os +import stat +import sys + +path = sys.argv[1] +max_bytes = int(sys.argv[2]) +# DIFY_WORKSPACE_CHECKPOINT: arguments_loaded +file_fd = None +try: + file_fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) + # DIFY_WORKSPACE_CHECKPOINT: file_opened + file_stat = os.fstat(file_fd) + if not stat.S_ISREG(file_stat.st_mode): + raise FileNotFoundError(path) + size = int(file_stat.st_size) + # DIFY_WORKSPACE_CHECKPOINT: file_size_captured + content = None + if size <= max_bytes: + chunks = [] + remaining = max_bytes + 1 + while remaining > 0: + chunk = os.read(file_fd, min(1024 * 1024, remaining)) + if not chunk: + break + chunks.append(chunk) + remaining -= len(chunk) + captured = b"".join(chunks) + if len(captured) > max_bytes: + size = max(size, len(captured)) + else: + content = captured +finally: + if file_fd is not None: + os.close(file_fd) + +if content is None: + payload = {"path": path, "size": size, "too_large": True} +else: + payload = { + "path": path, + "size": size, + "content_base64": base64.b64encode(content).decode("ascii"), + } +blob = base64.b64encode(json.dumps(payload).encode()).decode() +print("<<>>" + blob + "<<>>") +""" class ShellFileTransferError(RuntimeError): @@ -119,6 +265,8 @@ type ShellctlClientFactory = Callable[[], ShellctlClientProtocol] @dataclass(slots=True) class ShellctlCommands(ShellCommandProtocol): client: ShellctlClientProtocol + home_dir: str | None = None + workspace_dir: str | None = None async def run( self, @@ -128,7 +276,15 @@ class ShellctlCommands(ShellCommandProtocol): env: dict[str, str] | None = None, timeout: float, ) -> ShellCommandResult: - return _from_job_result(await _run_client_call(self.client.run(script, cwd=cwd, env=env, timeout=timeout))) + resolved_cwd = _resolve_lease_cwd( + cwd, + home_dir=self.home_dir, + workspace_dir=self.workspace_dir, + ) + resolved_env = _lease_env(env, home_dir=self.home_dir) + return _from_job_result( + await _run_client_call(self.client.run(script, cwd=resolved_cwd, env=resolved_env, timeout=timeout)) + ) async def wait( self, @@ -188,14 +344,127 @@ class ShellctlCommands(ShellCommandProtocol): @dataclass(slots=True) class ShellctlFileTransfer(ShellFileTransferProtocol): client: ShellctlClientProtocol + cwd: str | None = None + home_dir: str | None = None timeout: float = _FILE_TRANSFER_TIMEOUT_SECONDS + async def list_directory( + self, + *, + path: str, + limit: int, + ) -> WorkspaceListResult: + payload = await self._run_workspace_script( + _LIST_WORKSPACE_SCRIPT, + args=[self._resolve_path(path), str(limit)], + ) + raw_entries = payload.get("entries") + if not isinstance(raw_entries, list): + raise WorkspaceUnavailableError("workspace list returned an invalid entries payload") + entries: list[WorkspaceFileEntry] = [] + for raw_entry in raw_entries: + if not isinstance(raw_entry, dict): + raise WorkspaceUnavailableError("workspace list returned an invalid entry") + entries.append( + WorkspaceFileEntry( + name=str(raw_entry.get("name", "")), + type=str(raw_entry.get("type", "other")), + size=int(raw_entry["size"]) if isinstance(raw_entry.get("size"), int) else None, + mtime=int(raw_entry["mtime"]) if isinstance(raw_entry.get("mtime"), int) else None, + ) + ) + return WorkspaceListResult( + path=path, + entries=tuple(entries), + truncated=payload.get("truncated") is True, + ) + + async def read_file( + self, + *, + path: str, + max_bytes: int, + ) -> WorkspaceReadResult: + payload = await self._run_workspace_script( + _READ_WORKSPACE_SCRIPT, + args=[self._resolve_path(path), str(max_bytes)], + ) + size = payload.get("size") + if not isinstance(size, int): + raise WorkspaceUnavailableError("workspace read returned an invalid size") + text = payload.get("text") + return WorkspaceReadResult( + path=path, + size=size, + truncated=payload.get("truncated") is True, + binary=payload.get("binary") is True, + text=text if isinstance(text, str) else None, + ) + + async def read_bytes( + self, + *, + path: str, + max_bytes: int, + ) -> WorkspaceFileContent: + if max_bytes < 1: + raise ValueError("max_bytes must be positive") + payload = await self._run_workspace_script( + _READ_WORKSPACE_BYTES_SCRIPT, + args=[self._resolve_path(path), str(max_bytes)], + ) + size = payload.get("size") + if payload.get("too_large") is True: + if not isinstance(size, int): + raise WorkspaceUnavailableError("workspace bytes read returned an invalid size") + raise WorkspaceFileTooLargeError(path=path, size=size, max_bytes=max_bytes) + encoded = payload.get("content_base64") + if not isinstance(size, int) or not isinstance(encoded, str): + raise WorkspaceUnavailableError("workspace bytes read returned an invalid payload") + try: + content = base64.b64decode(encoded, validate=True) + except (ValueError, binascii.Error) as exc: + raise WorkspaceUnavailableError("workspace bytes read returned invalid base64") from exc + if len(content) != size: + raise WorkspaceUnavailableError("workspace bytes read returned an invalid size") + return WorkspaceFileContent( + path=path, + size=size, + content=content, + ) + + async def _run_workspace_script(self, source: str, *, args: list[str]) -> dict[str, object]: + command = _python_stdin_command(source, args=args) + completed = await _run_to_completion( + self.client, + command, + cwd=_resolve_lease_cwd(self.cwd, home_dir=self.home_dir, workspace_dir=self.cwd), + env=_lease_env(None, home_dir=self.home_dir), + timeout=self.timeout, + ) + if completed.exit_code != 0: + detail = _output_tail(completed.output) + if "PermissionError" in completed.output: + raise WorkspacePathError(detail) + raise WorkspaceUnavailableError(detail) + return _decode_workspace_payload(completed.output) + + def _resolve_path(self, path: str) -> str: + if self.home_dir is None: + return path + if path == "~": + return self.home_dir + if path.startswith("~/"): + return f"{self.home_dir.rstrip('/')}/{path[2:]}" + return path + async def upload(self, *, content: bytes, remote_path: str, cwd: str | None = None) -> None: encoded = base64.b64encode(content).decode("ascii") completed = await _run_to_completion( self.client, _upload_script(remote_path=remote_path, encoded=encoded), - cwd=cwd, + cwd=_resolve_lease_cwd(cwd, home_dir=self.home_dir, workspace_dir=self.cwd), + env=_lease_env(None, home_dir=self.home_dir), timeout=self.timeout, ) if completed.exit_code != 0: @@ -208,7 +477,8 @@ class ShellctlFileTransfer(ShellFileTransferProtocol): completed = await _run_to_completion( self.client, _download_script(remote_path=remote_path), - cwd=cwd, + cwd=_resolve_lease_cwd(cwd, home_dir=self.home_dir, workspace_dir=self.cwd), + env=_lease_env(None, home_dir=self.home_dir), timeout=self.timeout, ) if completed.exit_code == _DOWNLOAD_MISSING_EXIT_CODE: @@ -225,79 +495,6 @@ class ShellctlFileTransfer(ShellFileTransferProtocol): raise ShellFileTransferError(f"Downloaded payload for {remote_path!r} was not valid base64.") from exc -@dataclass(slots=True) -class ShellctlResource(ShellResourceProtocol): - """Live shellctl connection. - - For shellctl there is no separate sandbox lifecycle: the shellctl server is a - long-running process that persists across runs. The ``sandbox_id`` is the - shellctl entrypoint URL, used as a stable identifier so the layer can - distinguish ``create()`` (first run) from ``attach()`` (subsequent runs). - Both ``suspend()`` and ``delete()`` simply close the HTTP client; the server - and its filesystem remain intact either way. - """ - - client: ShellctlClientProtocol - _commands: ShellCommandProtocol - _files: ShellFileTransferProtocol - _sandbox_id: str | None = None - - @property - def commands(self) -> ShellCommandProtocol: - return self._commands - - @property - def files(self) -> ShellFileTransferProtocol: - return self._files - - @property - def sandbox_id(self) -> str | None: - return self._sandbox_id - - async def suspend(self) -> None: - await self._close_client() - - async def delete(self) -> None: - await self._close_client() - - async def _close_client(self) -> None: - try: - await self.client.close() - except RuntimeError as exc: - raise _map_error(exc) from exc - - -@dataclass(slots=True) -class ShellctlProvider(ShellProviderProtocol): - entrypoint: str - token: str - output_limit: int = _SHELLCTL_OUTPUT_LIMIT_BYTES - client_factory: ShellctlClientFactory | None = None - - async def create(self) -> ShellctlResource: - return self._build_resource(sandbox_id=self.entrypoint) - - async def attach(self, sandbox_id: str) -> ShellctlResource: - return self._build_resource(sandbox_id=sandbox_id) - - def _build_resource(self, *, sandbox_id: str | None) -> ShellctlResource: - client = ( - self.client_factory() - if self.client_factory is not None - else create_default_shellctl_client_factory( - entrypoint=self.entrypoint, - token=self.token, - output_limit=self.output_limit, - )() - ) - return ShellctlResource( - client=client, - _commands=ShellctlCommands(client=client), - _files=ShellctlFileTransfer(client=client), - _sandbox_id=sandbox_id, - ) - - def create_default_shellctl_client_factory( *, entrypoint: str, @@ -334,12 +531,16 @@ async def _run_client_call(awaitable: Awaitable[ResultT]) -> ResultT: raise ShellProviderError(str(exc), code="timeout") from exc except httpx.RequestError as exc: raise ShellProviderError(str(exc), code="request_error") from exc - except RuntimeError as exc: + except ShellctlClientError as exc: raise _map_error(exc) from exc -def _map_error(exc: RuntimeError) -> ShellProviderError: - return ShellProviderError(str(exc), code=getattr(exc, "code", None)) +def _map_error(exc: ShellctlClientError) -> ShellProviderError: + return ShellProviderError( + str(exc), + code=exc.code, + status_code=exc.status_code, + ) def _from_job_result(result: ShellctlJobResult) -> ShellCommandResult: @@ -379,12 +580,13 @@ async def _run_to_completion( script: str, *, cwd: str | None, + env: dict[str, str] | None, timeout: float, ) -> _CompletedShellctlJob: deadline = time.monotonic() + timeout job_id: str | None = None try: - result = await _run_client_call(client.run(script, cwd=cwd, env=None, timeout=_remaining_timeout(deadline))) + result = await _run_client_call(client.run(script, cwd=cwd, env=env, timeout=_remaining_timeout(deadline))) parts = [result.output] job_id = result.job_id while not result.done or result.truncated: @@ -422,6 +624,26 @@ def _download_script(*, remote_path: str) -> str: ) +def _python_stdin_command(source: str, *, args: list[str]) -> str: + quoted_args = " ".join(shlex.quote(value) for value in args) + return f"python3 - {quoted_args} <<'PY'\n{source.strip()}\nPY" + + +def _decode_workspace_payload(output: str) -> dict[str, object]: + begin = output.find(_WORKSPACE_PAYLOAD_BEGIN) + end = output.find(_WORKSPACE_PAYLOAD_END, begin + len(_WORKSPACE_PAYLOAD_BEGIN)) if begin >= 0 else -1 + if begin < 0 or end < 0: + raise WorkspaceUnavailableError("workspace command returned no framed payload") + encoded = "".join(output[begin + len(_WORKSPACE_PAYLOAD_BEGIN) : end].split()) + try: + value = json.loads(base64.b64decode(encoded, validate=True).decode("utf-8")) + except (binascii.Error, UnicodeDecodeError, ValueError) as exc: + raise WorkspaceUnavailableError("workspace command returned an invalid framed payload") from exc + if not isinstance(value, dict): + raise WorkspaceUnavailableError("workspace command returned a non-object payload") + return cast(dict[str, object], value) + + def _extract_transfer_payload(output: str) -> str: pattern = re.escape(_TRANSFER_BEGIN) + r"(.*?)" + re.escape(_TRANSFER_END) match = re.search(pattern, output, re.DOTALL) @@ -441,6 +663,39 @@ def _remaining_timeout(deadline: float) -> float: return remaining +def _lease_env(env: dict[str, str] | None, *, home_dir: str | None) -> dict[str, str] | None: + if home_dir is None: + return env + resolved = dict(env or {}) + resolved["HOME"] = home_dir + return resolved + + +def _resolve_lease_cwd( + cwd: str | None, + *, + home_dir: str | None, + workspace_dir: str | None, +) -> str | None: + if home_dir is None or workspace_dir is None: + return cwd + if cwd is None or cwd == "": + candidate = workspace_dir + elif cwd == "~": + candidate = home_dir + elif cwd.startswith("~/"): + candidate = posixpath.join(home_dir, cwd[2:]) + elif posixpath.isabs(cwd): + candidate = cwd + else: + candidate = posixpath.join(workspace_dir, cwd) + candidate = posixpath.normpath(candidate) + roots = (posixpath.normpath(home_dir), posixpath.normpath(workspace_dir)) + if not any(posixpath.commonpath((candidate, root)) == root for root in roots): + raise ValueError("shell cwd is outside this RuntimeLease Home and Workspace") + return candidate + + def _shquote(value: str) -> str: return "'" + value.replace("'", "'\\''") + "'" @@ -451,7 +706,5 @@ __all__ = [ "ShellctlClientProtocol", "ShellctlCommands", "ShellctlFileTransfer", - "ShellctlProvider", - "ShellctlResource", "create_default_shellctl_client_factory", ] diff --git a/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py b/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py index f3efe34f00a..59cd43b6bb5 100644 --- a/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py +++ b/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py @@ -15,6 +15,7 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Protocol +from urllib.parse import urljoin import httpx from pydantic import BaseModel, ConfigDict, ValidationError @@ -132,7 +133,9 @@ class DifyApiAgentStubFileRequestHandler: upload_url = data.get("url") if not isinstance(upload_url, str) or not upload_url: raise AgentStubFileRequestError(502, "Dify API upload request response is missing url") - return AgentStubFileUploadResponse(upload_url=upload_url) + return AgentStubFileUploadResponse( + upload_url=urljoin(f"{self.inner_api_url.rstrip('/')}/", upload_url), + ) async def create_download_request( self, @@ -154,7 +157,7 @@ class DifyApiAgentStubFileRequestHandler: not match ``AgentStubFileDownloadResponse``. """ execution_context = self._require_user_context(principal.execution_context) - payload = { + payload: dict[str, object] = { "tenant_id": execution_context.tenant_id, "user_id": execution_context.user_id, "user_from": execution_context.user_from, diff --git a/dify-agent/src/dify_agent/client/_client.py b/dify-agent/src/dify_agent/client/_client.py index 9c15d96cb59..cf791e52e27 100644 --- a/dify-agent/src/dify_agent/client/_client.py +++ b/dify-agent/src/dify_agent/client/_client.py @@ -1,7 +1,7 @@ """HTTPX-based client for the Dify Agent HTTP API. The client uses the public DTOs from ``dify_agent.protocol`` for request and -response parsing across both run-management and sandbox-file endpoints. It +response parsing across run-management and working-environment endpoints. It intentionally does not retry non-idempotent ``POST`` requests such as ``/runs``. SSE streams are the only operation with reconnect logic: transient stream, connect, or read failures, stream timeouts, and HTTP 5xx stream @@ -30,17 +30,23 @@ from dify_agent.protocol import ( CancelRunResponse, CreateRunRequest, CreateRunResponse, + CreateExecutionBindingRequest, + CreateExecutionBindingResponse, + CreateHomeSnapshotFromBindingRequest, + DeleteHomeSnapshotRequest, + DestroyExecutionBindingRequest, + HomeSnapshotResponse, + InitializeHomeSnapshotRequest, RUN_EVENT_ADAPTER, RunEvent, RunEventsResponse, RunStatusResponse, - SandboxListRequest, - SandboxListResponse, - SandboxLocator, - SandboxReadRequest, - SandboxReadResponse, - SandboxUploadRequest, - SandboxUploadResponse, + WorkspaceListRequest, + WorkspaceListResponse, + WorkspaceReadRequest, + WorkspaceReadResponse, + WorkspaceUploadRequest, + WorkspaceUploadResponse, ) _ResponseModelT = TypeVar("_ResponseModelT", bound=BaseModel) @@ -469,61 +475,131 @@ class Client: raise DifyAgentClientError(f"get_events_sync request failed: {exc}") from exc return _parse_model_response(response, RunEventsResponse) - async def list_sandbox_files(self, locator: SandboxLocator, path: str) -> SandboxListResponse: - """List a sandbox directory through ``POST /sandbox/files/list``.""" - request_model = _build_request_model(SandboxListRequest, locator=locator, path=path) - response = await self._post_async_json("list_sandbox_files", "/sandbox/files/list", request_model) - return _parse_model_response(response, SandboxListResponse) + async def list_workspace_files(self, backend_binding_ref: str, path: str) -> WorkspaceListResponse: + request_model = WorkspaceListRequest(backend_binding_ref=backend_binding_ref, path=path) + response = await self._post_async_json("list_workspace_files", "/workspace/files/list", request_model) + return _parse_model_response(response, WorkspaceListResponse) - def list_sandbox_files_sync(self, locator: SandboxLocator, path: str) -> SandboxListResponse: - """Synchronous variant of ``list_sandbox_files``.""" - request_model = _build_request_model(SandboxListRequest, locator=locator, path=path) - response = self._post_sync_json("list_sandbox_files_sync", "/sandbox/files/list", request_model) - return _parse_model_response(response, SandboxListResponse) + def list_workspace_files_sync(self, backend_binding_ref: str, path: str) -> WorkspaceListResponse: + request_model = WorkspaceListRequest(backend_binding_ref=backend_binding_ref, path=path) + response = self._post_sync_json("list_workspace_files_sync", "/workspace/files/list", request_model) + return _parse_model_response(response, WorkspaceListResponse) - async def read_sandbox_file( + async def read_workspace_file( self, - locator: SandboxLocator, + backend_binding_ref: str, path: str, max_bytes: int = 262144, - ) -> SandboxReadResponse: - """Read a sandbox file preview through ``POST /sandbox/files/read``.""" - request_model = _build_request_model( - SandboxReadRequest, - locator=locator, - path=path, - max_bytes=max_bytes, - ) - response = await self._post_async_json("read_sandbox_file", "/sandbox/files/read", request_model) - return _parse_model_response(response, SandboxReadResponse) + ) -> WorkspaceReadResponse: + request_model = WorkspaceReadRequest(backend_binding_ref=backend_binding_ref, path=path, max_bytes=max_bytes) + response = await self._post_async_json("read_workspace_file", "/workspace/files/read", request_model) + return _parse_model_response(response, WorkspaceReadResponse) - def read_sandbox_file_sync( + def read_workspace_file_sync( self, - locator: SandboxLocator, + backend_binding_ref: str, path: str, max_bytes: int = 262144, - ) -> SandboxReadResponse: - """Synchronous variant of ``read_sandbox_file``.""" - request_model = _build_request_model( - SandboxReadRequest, - locator=locator, - path=path, - max_bytes=max_bytes, + ) -> WorkspaceReadResponse: + request_model = WorkspaceReadRequest(backend_binding_ref=backend_binding_ref, path=path, max_bytes=max_bytes) + response = self._post_sync_json("read_workspace_file_sync", "/workspace/files/read", request_model) + return _parse_model_response(response, WorkspaceReadResponse) + + async def upload_workspace_file(self, request: WorkspaceUploadRequest) -> WorkspaceUploadResponse: + response = await self._post_async_json("upload_workspace_file", "/workspace/files/upload", request) + return _parse_model_response(response, WorkspaceUploadResponse) + + def upload_workspace_file_sync(self, request: WorkspaceUploadRequest) -> WorkspaceUploadResponse: + response = self._post_sync_json("upload_workspace_file_sync", "/workspace/files/upload", request) + return _parse_model_response(response, WorkspaceUploadResponse) + + async def create_execution_binding(self, request: CreateExecutionBindingRequest) -> CreateExecutionBindingResponse: + response = await self._post_async_json("create_execution_binding", "/execution-bindings", request) + return _parse_model_response(response, CreateExecutionBindingResponse) + + def create_execution_binding_sync(self, request: CreateExecutionBindingRequest) -> CreateExecutionBindingResponse: + response = self._post_sync_json("create_execution_binding_sync", "/execution-bindings", request) + return _parse_model_response(response, CreateExecutionBindingResponse) + + async def destroy_execution_binding(self, request: DestroyExecutionBindingRequest) -> None: + response = await self._post_async_json("destroy_execution_binding", "/execution-bindings/destroy", request) + _raise_for_status(response) + + def destroy_execution_binding_sync(self, request: DestroyExecutionBindingRequest) -> None: + response = self._post_sync_json("destroy_execution_binding_sync", "/execution-bindings/destroy", request) + _raise_for_status(response) + + async def initialize_home_snapshot(self, request: InitializeHomeSnapshotRequest) -> HomeSnapshotResponse: + """Create a backend-native initial Home Snapshot.""" + response = await self._post_async_json( + "initialize_home_snapshot", + "/home-snapshots/initialize", + request, ) - response = self._post_sync_json("read_sandbox_file_sync", "/sandbox/files/read", request_model) - return _parse_model_response(response, SandboxReadResponse) + return _parse_model_response(response, HomeSnapshotResponse) - async def upload_sandbox_file(self, locator: SandboxLocator, path: str) -> SandboxUploadResponse: - """Upload a sandbox file mapping through ``POST /sandbox/files/upload``.""" - request_model = _build_request_model(SandboxUploadRequest, locator=locator, path=path) - response = await self._post_async_json("upload_sandbox_file", "/sandbox/files/upload", request_model) - return _parse_model_response(response, SandboxUploadResponse) + def initialize_home_snapshot_sync(self, request: InitializeHomeSnapshotRequest) -> HomeSnapshotResponse: + """Synchronous variant of ``initialize_home_snapshot``.""" + response = self._post_sync_json( + "initialize_home_snapshot_sync", + "/home-snapshots/initialize", + request, + ) + return _parse_model_response(response, HomeSnapshotResponse) - def upload_sandbox_file_sync(self, locator: SandboxLocator, path: str) -> SandboxUploadResponse: - """Synchronous variant of ``upload_sandbox_file``.""" - request_model = _build_request_model(SandboxUploadRequest, locator=locator, path=path) - response = self._post_sync_json("upload_sandbox_file_sync", "/sandbox/files/upload", request_model) - return _parse_model_response(response, SandboxUploadResponse) + async def create_home_snapshot_from_binding( + self, + request: CreateHomeSnapshotFromBindingRequest, + ) -> HomeSnapshotResponse: + """Checkpoint Home from the exact Execution Binding identified by the request.""" + response = await self._post_async_json( + "create_home_snapshot_from_binding", + "/home-snapshots/from-binding", + request, + ) + return _parse_model_response(response, HomeSnapshotResponse) + + def create_home_snapshot_from_binding_sync( + self, + request: CreateHomeSnapshotFromBindingRequest, + ) -> HomeSnapshotResponse: + """Synchronous variant of ``create_home_snapshot_from_binding``.""" + response = self._post_sync_json( + "create_home_snapshot_from_binding_sync", + "/home-snapshots/from-binding", + request, + ) + return _parse_model_response(response, HomeSnapshotResponse) + + async def delete_home_snapshot(self, snapshot_ref: str) -> None: + """Idempotently delete one backend Home Snapshot.""" + try: + response = await self._get_async_http_client().post( + self._url("/home-snapshots/delete"), + content=DeleteHomeSnapshotRequest(snapshot_ref=snapshot_ref).model_dump_json(), + headers=self._merged_headers({"Content-Type": "application/json"}), + timeout=self._timeout, + ) + except httpx.TimeoutException as exc: + raise DifyAgentTimeoutError("delete_home_snapshot timed out") from exc + except httpx.RequestError as exc: + raise DifyAgentClientError(f"delete_home_snapshot request failed: {exc}") from exc + _raise_for_status(response) + + def delete_home_snapshot_sync(self, snapshot_ref: str) -> None: + """Synchronous variant of ``delete_home_snapshot``.""" + try: + response = self._get_sync_http_client().post( + self._url("/home-snapshots/delete"), + content=DeleteHomeSnapshotRequest(snapshot_ref=snapshot_ref).model_dump_json(), + headers=self._merged_headers({"Content-Type": "application/json"}), + timeout=self._timeout, + ) + except httpx.TimeoutException as exc: + raise DifyAgentTimeoutError("delete_home_snapshot_sync timed out") from exc + except httpx.RequestError as exc: + raise DifyAgentClientError(f"delete_home_snapshot_sync request failed: {exc}") from exc + _raise_for_status(response) async def stream_events( self, diff --git a/dify-agent/src/dify_agent/layers/runtime/__init__.py b/dify-agent/src/dify_agent/layers/runtime/__init__.py new file mode 100644 index 00000000000..872ed9120c4 --- /dev/null +++ b/dify-agent/src/dify_agent/layers/runtime/__init__.py @@ -0,0 +1,4 @@ +from .configs import DIFY_RUNTIME_LAYER_TYPE_ID, DifyRuntimeLayerConfig +from .layer import DifyRuntimeLayer + +__all__ = ["DIFY_RUNTIME_LAYER_TYPE_ID", "DifyRuntimeLayer", "DifyRuntimeLayerConfig"] diff --git a/dify-agent/src/dify_agent/layers/runtime/configs.py b/dify-agent/src/dify_agent/layers/runtime/configs.py new file mode 100644 index 00000000000..f7f3a576dbb --- /dev/null +++ b/dify-agent/src/dify_agent/layers/runtime/configs.py @@ -0,0 +1,17 @@ +"""Client-safe config for operation-scoped RuntimeLease acquisition.""" + +from typing import ClassVar, Final + +from agenton.layers import LayerConfig +from pydantic import ConfigDict, Field + +DIFY_RUNTIME_LAYER_TYPE_ID: Final[str] = "dify.runtime" + + +class DifyRuntimeLayerConfig(LayerConfig): + backend_binding_ref: str = Field(min_length=1) + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +__all__ = ["DIFY_RUNTIME_LAYER_TYPE_ID", "DifyRuntimeLayerConfig"] diff --git a/dify-agent/src/dify_agent/layers/runtime/layer.py b/dify-agent/src/dify_agent/layers/runtime/layer.py new file mode 100644 index 00000000000..e6fbaf70aff --- /dev/null +++ b/dify-agent/src/dify_agent/layers/runtime/layer.py @@ -0,0 +1,75 @@ +"""Agenton layer exposing one operation-scoped RuntimeLease.""" + +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from typing import ClassVar + +from agenton.layers import EmptyRuntimeState, NoLayerDeps, PlainLayer +from typing_extensions import Self, override + +from dify_agent.layers.runtime.configs import DIFY_RUNTIME_LAYER_TYPE_ID, DifyRuntimeLayerConfig +from dify_agent.runtime_backend import ExecutionBindingBackend, RuntimeLease +from dify_agent.runtime_backend.leases import open_runtime_lease + + +@dataclass(slots=True) +class DifyRuntimeLayer(PlainLayer[NoLayerDeps, DifyRuntimeLayerConfig, EmptyRuntimeState]): + """Acquire/release a Binding without owning its product lifecycle.""" + + type_id: ClassVar[str | None] = DIFY_RUNTIME_LAYER_TYPE_ID + config: DifyRuntimeLayerConfig + backend: ExecutionBindingBackend + _lease: RuntimeLease | None = field(default=None, init=False) + + @classmethod + @override + def from_config(cls, config: DifyRuntimeLayerConfig) -> Self: + del config + raise TypeError("DifyRuntimeLayer requires a server-injected ExecutionBindingBackend") + + @classmethod + def from_config_with_backend( + cls, + config: DifyRuntimeLayerConfig, + *, + backend: ExecutionBindingBackend, + ) -> Self: + return cls(config=DifyRuntimeLayerConfig.model_validate(config), backend=backend) + + @property + def lease(self) -> RuntimeLease: + if self._lease is None: + raise RuntimeError("DifyRuntimeLayer lease is only available inside resource_context()") + return self._lease + + @override + @asynccontextmanager + async def resource_context(self) -> AsyncGenerator[None]: + if self._lease is not None: + raise RuntimeError("DifyRuntimeLayer resource_context() is already active") + async with open_runtime_lease(self.backend, self.config.backend_binding_ref) as lease: + self._lease = lease + try: + yield + finally: + self._lease = None + + @override + async def on_context_create(self) -> None: + _ = self.lease + + @override + async def on_context_resume(self) -> None: + _ = self.lease + + @override + async def on_context_suspend(self) -> None: + _ = self.lease + + @override + async def on_context_delete(self) -> None: + _ = self.lease + + +__all__ = ["DifyRuntimeLayer"] diff --git a/dify-agent/src/dify_agent/layers/shell/__init__.py b/dify-agent/src/dify_agent/layers/shell/__init__.py index 69af3150c76..34800d69471 100644 --- a/dify-agent/src/dify_agent/layers/shell/__init__.py +++ b/dify-agent/src/dify_agent/layers/shell/__init__.py @@ -1,8 +1,8 @@ """Client-safe exports for the Dify shell layer DTOs. -The runtime layer implementation lives in ``layer.py`` and imports shellctl -client code plus server-side lifecycle behavior. Keep this package root -import-safe for client code that only needs to build run requests. +The runtime layer implementation lives in ``layer.py`` and consumes a +server-injected active Sandbox lease. Keep this package root import-safe for +client code that only needs to build run requests. """ from dify_agent.layers.shell.configs import ( @@ -10,7 +10,6 @@ from dify_agent.layers.shell.configs import ( DifyShellCliToolConfig, DifyShellEnvVarConfig, DifyShellLayerConfig, - DifyShellSandboxConfig, DifyShellSecretRefConfig, ) @@ -19,6 +18,5 @@ __all__ = [ "DifyShellCliToolConfig", "DifyShellEnvVarConfig", "DifyShellLayerConfig", - "DifyShellSandboxConfig", "DifyShellSecretRefConfig", ] diff --git a/dify-agent/src/dify_agent/layers/shell/configs.py b/dify-agent/src/dify_agent/layers/shell/configs.py index 63f1faf5d5e..821e02cc89c 100644 --- a/dify-agent/src/dify_agent/layers/shell/configs.py +++ b/dify-agent/src/dify_agent/layers/shell/configs.py @@ -1,10 +1,11 @@ """Client-safe DTOs for the Dify shell Agenton layer. -Server-only shellctl connection settings are injected by the runtime provider -factory. Public config carries product-level Agent Soul settings that must affect -the sandbox workspace itself: CLI tool bootstrap commands, normal environment -variables, secret environment variable names, sandbox-provider metadata, and the -Agent Stub drive ref used by shell-visible drive commands. +Server-only Agent Stub and redaction settings are injected by the runtime +provider factory. The Sandbox dependency supplies the active shellctl data +plane. Public config carries product-level Agent Soul settings that affect the +workspace itself: CLI tool bootstrap commands, normal environment variables, +secret environment variable names, and the Agent Stub drive ref used by +shell-visible drive commands. Sandbox selection is a deployment concern. """ import re @@ -67,24 +68,8 @@ class DifyShellCliToolConfig(BaseModel): return [command for command in (item.strip() for item in value) if command] -class DifyShellSandboxConfig(BaseModel): - """Sandbox provider selection persisted in Agent Soul.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - provider: str | None = Field(default=None, max_length=255) - config: dict[str, object] = Field(default_factory=dict) - - -class DifyShellEnterpriseSandboxConfig(DifyShellSandboxConfig): - """Enterprise sandbox provider configuration.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - gateway_endpoint: str = Field(..., max_length=255) - - class DifyShellLayerConfig(LayerConfig): - """Public config for the shellctl-backed Dify shell layer.""" + """Public product behavior for the Sandbox-backed Dify Shell layer.""" model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") @@ -93,7 +78,6 @@ class DifyShellLayerConfig(LayerConfig): cli_tools: list[DifyShellCliToolConfig] = Field(default_factory=list) env: list[DifyShellEnvVarConfig] = Field(default_factory=list) secret_refs: list[DifyShellSecretRefConfig] = Field(default_factory=list) - sandbox: DifyShellSandboxConfig | None = None redact_patterns: list[str] = Field(default_factory=list) @@ -102,6 +86,5 @@ __all__ = [ "DifyShellCliToolConfig", "DifyShellEnvVarConfig", "DifyShellLayerConfig", - "DifyShellSandboxConfig", "DifyShellSecretRefConfig", ] diff --git a/dify-agent/src/dify_agent/layers/shell/layer.py b/dify-agent/src/dify_agent/layers/shell/layer.py index b38a151ddee..f987ad07371 100644 --- a/dify-agent/src/dify_agent/layers/shell/layer.py +++ b/dify-agent/src/dify_agent/layers/shell/layer.py @@ -1,37 +1,11 @@ -"""Shell runtime layer backed by a live shell provider resource. - -Shell command execution requires a bound execution-context layer with a safe -``agent_id``. The layer uses the current bound execution context to run -commands with ``HOME=/`` and a home-rooted workspace path. The -persisted runtime state intentionally keeps the historical -``~/workspace/`` identity so existing session snapshots stay -compatible while live command execution no longer depends on the sandbox user's -ambient home directory. Entering or re-entering the layer re-ensures the live -home/workspace directories for the currently bound ``agent_id`` before user -commands are sent. - -Sandbox lifecycle: - The shell provider exposes four operations: ``create``, ``attach``, - ``suspend``, and ``delete``. On the first run (no ``sandbox_id`` in - runtime state) ``resource_context()`` calls ``create()`` to provision a - new sandbox and persists the returned ``sandbox_id``. On subsequent runs - it calls ``attach(sandbox_id)`` to re-connect to the existing sandbox. - If the sandbox has expired (the provider raises ``SandboxExpiredError``), - the error propagates to the caller — the user must start a new session. - On normal exit (suspend) the resource is detached via ``suspend()``, - keeping the sandbox alive. On final cleanup (``on_context_delete``) the - resource is destroyed via ``delete()``. This allows the enterprise - provider to reuse the same sandbox pod across conversation turns. -""" +"""Shell tools over the data plane exposed by the active Runtime layer.""" from __future__ import annotations -from collections.abc import AsyncGenerator, Sequence -from contextlib import asynccontextmanager +from collections.abc import Sequence import json import logging import re -import secrets import time from dataclasses import dataclass, field from typing import ClassVar, Literal, NotRequired, Protocol, TypedDict, runtime_checkable @@ -54,14 +28,14 @@ from dify_agent.adapters.shell.protocols import ( ShellCommandProtocol, ShellCommandResult, ShellPromptObservation, - ShellProviderProtocol, - ShellResourceProtocol, ) from dify_agent.agent_stub.protocol import AGENT_STUB_AUTH_JWE_ENV_VAR from dify_agent.agent_stub.shell_env import ShellAgentStubTokenFactory, build_shell_agent_stub_env from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig +from dify_agent.layers.runtime.layer import DifyRuntimeLayer from dify_agent.layers.shell.configs import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig from dify_agent.layers.shell.output_text import normalized_output_text, utf8_prefix, utf8_suffix +from dify_agent.runtime_backend import RuntimeLease logger = logging.getLogger(__name__) @@ -74,14 +48,6 @@ class _HasErrorCode(Protocol): DEFAULT_TIMEOUT_SECONDS = 30.0 DEFAULT_TERMINATE_GRACE_SECONDS = 10.0 -_WORKSPACE_ROOT = "~/workspace" -_WORKSPACE_DIR_NAME = "workspace" -_WORKSPACE_COLLISION_EXIT_CODE = 17 -_SESSION_TIME_HEX_MASK = 0xFFFFF -_SESSION_RANDOM_HEX_LENGTH = 2 -_SESSION_ID_ATTEMPT_LIMIT = 256 -_SESSION_ID_PATTERN = re.compile(r"^[0-9a-f]{7}$") -_AGENT_HOME_SEGMENT_PATTERN = re.compile(r"^[A-Za-z0-9._-]+$") _SHELL_OUTPUT_PROMPT_EDGE_BYTES = 8 * 1024 _SHELLCTL_OUTPUT_LIMIT_BYTES = 2 * _SHELL_OUTPUT_PROMPT_EDGE_BYTES _REMOTE_COMPLETE_OUTPUT_MAX_BYTES = 1024 * 1024 @@ -142,8 +108,8 @@ Installed CLI: Workspace persistence rules: -- The current workspace cwd is stable during this run, but it is temporary and may be deleted later. -- Do not treat files in the current workspace cwd as persisted state. +- The current workspace cwd is stable across runs for the current product session. +- Workspace files are working data, not Agent configuration, and are removed when that product session ends. - In build mode, config changes persist only after you run the matching `dify-agent config ...` mutation command. - Shell file edits alone do not save Agent config files, skills, env, or notes. - In non-build modes, local shell changes are not a persistence mechanism for Agent configuration. @@ -202,24 +168,15 @@ type ShellInterruptToolResult = str | ShellToolErrorObservation class DifyShellLayerDeps(LayerDeps): execution_context: PlainLayer[NoLayerDeps, DifyExecutionContextLayerConfig, EmptyRuntimeState] | None # pyright: ignore[reportUninitializedInstanceVariable] + runtime: DifyRuntimeLayer # pyright: ignore[reportUninitializedInstanceVariable] class DifyShellRuntimeState(BaseModel): - session_id: str | None = None - workspace_cwd: str | None = None - sandbox_id: str | None = None job_ids: list[str] = Field(default_factory=list) job_offsets: dict[str, NonNegativeInt] = Field(default_factory=dict) model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid", validate_assignment=True) - @field_validator("session_id") - @classmethod - def validate_session_id(cls, value: str | None) -> str | None: - if value is None: - return value - return _validated_session_id(value) - @field_validator("job_ids") @classmethod def validate_job_ids(cls, value: list[str]) -> list[str]: @@ -228,13 +185,7 @@ class DifyShellRuntimeState(BaseModel): return value @model_validator(mode="after") - def validate_workspace_and_offsets(self) -> Self: - if self.workspace_cwd is not None: - if self.session_id is None: - raise ValueError("workspace_cwd requires a matching session_id.") - expected_workspace = _workspace_cwd(self.session_id) - if self.workspace_cwd != expected_workspace: - raise ValueError(f"workspace_cwd must equal {expected_workspace!r} for session_id {self.session_id!r}.") + def validate_job_offsets(self) -> Self: unknown_offset_job_ids = set(self.job_offsets) - set(self.job_ids) if unknown_offset_job_ids: names = ", ".join(sorted(unknown_offset_job_ids)) @@ -247,46 +198,43 @@ CompleteRemoteCommandResult = CompleteShellCommandResult @dataclass(slots=True) class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerConfig, DifyShellRuntimeState]): + """Expose Shell tools over the active RuntimeLease without owning it. + + Create optionally bootstraps configured CLI tools in the lease's Workspace. + Suspend and delete best-effort remove tracked shellctl jobs, then clear job + ids and offsets so they do not persist across requests. Commands, files, + Home, and cwd come only from ``DifyRuntimeLayer.lease``. Persistent Binding + and Workspace lifecycle remains exclusively owned by Dify API. + """ + type_id: ClassVar[str | None] = DIFY_SHELL_LAYER_TYPE_ID config: DifyShellLayerConfig - shell_provider: ShellProviderProtocol - shell_home_root: str = "/home" shell_redact_patterns: list[str] = field(default_factory=list) agent_stub_api_base_url: str | None = None agent_stub_token_factory: ShellAgentStubTokenFactory | None = None - _shell_resource: ShellResourceProtocol | None = None - _resource_should_delete: bool = False @classmethod @override def from_config(cls, config: DifyShellLayerConfig) -> Self: del config - raise TypeError("DifyShellLayer requires a shell provider and must use a provider factory.") + raise TypeError("DifyShellLayer requires server-injected settings and must use a provider factory.") @classmethod def from_config_with_settings( cls, config: DifyShellLayerConfig, *, - shell_provider: ShellProviderProtocol | None, - shell_home_root: str = "/home", shell_redact_patterns: list[str] | None = None, agent_stub_api_base_url: str | None = None, agent_stub_token_factory: ShellAgentStubTokenFactory | None = None, ) -> Self: - if shell_provider is None: - raise ValueError("DifyShellLayer requires a non-null shell provider when the 'dify.shell' layer is used.") - layer = cls( + return cls( config=config, - shell_provider=shell_provider, - shell_home_root=_normalize_shell_home_root(shell_home_root), shell_redact_patterns=shell_redact_patterns or [], agent_stub_api_base_url=agent_stub_api_base_url, agent_stub_token_factory=agent_stub_token_factory, ) - layer.bind_deps({}) - return layer @property @override @@ -308,95 +256,31 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC Tool(self._tool_interrupt, name="shell_interrupt"), ] - @override - @asynccontextmanager - async def resource_context(self) -> AsyncGenerator[None]: - """Acquire the live shell resource for one run invocation. - - On the first run (no ``sandbox_id`` in runtime state) the provider's - ``create()`` is called to provision a new sandbox. On subsequent runs - ``attach(sandbox_id)`` re-connects to the existing sandbox. If the - sandbox has expired (``SandboxExpiredError``), the error propagates - to the caller — the user must start a new session. - - On exit, ``suspend()`` is called by default to keep the sandbox alive. - If ``on_context_delete()`` ran (setting ``_resource_should_delete``), - ``delete()`` is called instead to destroy the sandbox. - """ - if self._shell_resource is not None: - raise RuntimeError("DifyShellLayer resource_context() is already active for this layer instance.") - sandbox_id = self.runtime_state.sandbox_id - if sandbox_id is not None: - resource = await self.shell_provider.attach(sandbox_id) - else: - resource = await self.shell_provider.create() - self.runtime_state = DifyShellRuntimeState.model_validate( - { - **self.runtime_state.model_dump(mode="python"), - "sandbox_id": resource.sandbox_id, - } - ) - self._shell_resource = resource - self._resource_should_delete = False - try: - yield - finally: - self._shell_resource = None - if self._resource_should_delete: - await resource.delete() - else: - await resource.suspend() - @override async def on_context_create(self) -> None: - _ = self._require_resource() - session_id: str | None = None - try: - session_id, workspace_cwd = await self._allocate_workspace() - await self._bootstrap_workspace(session_id) - except BaseException: - if session_id is not None: - await self._cleanup_workspace_best_effort(session_id) - self._resource_should_delete = True - raise - self.runtime_state = DifyShellRuntimeState.model_validate( - { - **self.runtime_state.model_dump(mode="python"), - "session_id": session_id, - "workspace_cwd": workspace_cwd, - } - ) + bootstrap_script = _workspace_bootstrap_script(self.config) + if not bootstrap_script: + return + result = await self._run_internal_script_complete(bootstrap_script, cwd=self._require_workspace_cwd()) + if result.exit_code != 0 or not result.output_complete: + raise RuntimeError( + f"Failed to bootstrap shell workspace {self._require_workspace_cwd()}: " + f"{result.status} exit_code={result.exit_code}" + ) @override async def on_context_resume(self) -> None: _ = self._require_resource() - session_id, _workspace_cwd = self._require_session_identity() - await self._ensure_live_workspace_exists(session_id) @override async def on_context_suspend(self) -> None: - _ = self._require_resource() + await self._delete_tracked_jobs_best_effort(self.runtime_state.job_ids) + self._clear_tracked_jobs() @override async def on_context_delete(self) -> None: - _ = self._require_resource() - identity = self._try_session_identity() - if identity is not None: - session_id, _workspace_cwd = identity - result = await self._run_internal_script_complete( - _workspace_cleanup_script(session_id=session_id), cwd=None - ) - if result.exit_code != 0 or not result.output_complete: - logger.warning( - "Shell workspace cleanup for session %s ended with status=%s exit_code=%s output_complete=%s.", - session_id, - result.status, - result.exit_code, - result.output_complete, - ) await self._delete_tracked_jobs_best_effort(self.runtime_state.job_ids) self._clear_tracked_jobs() - self._resource_should_delete = True async def _tool_run(self, script: str, timeout: float = DEFAULT_TIMEOUT_SECONDS) -> ShellRunToolResult: try: @@ -426,7 +310,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc) except Exception as exc: - return _tool_unexpected_error("shell_run", exc, session_id=self.runtime_state.session_id) + return _tool_unexpected_error("shell_run", exc) async def _tool_wait(self, job_id: str, timeout: float = DEFAULT_TIMEOUT_SECONDS) -> ShellRunToolResult: try: @@ -452,7 +336,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc, job_id=job_id) except Exception as exc: - return _tool_unexpected_error("shell_wait", exc, session_id=self.runtime_state.session_id, job_id=job_id) + return _tool_unexpected_error("shell_wait", exc, job_id=job_id) async def _tool_input(self, job_id: str, text: str, timeout: float = DEFAULT_TIMEOUT_SECONDS) -> ShellRunToolResult: try: @@ -478,7 +362,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc, job_id=job_id) except Exception as exc: - return _tool_unexpected_error("shell_input", exc, session_id=self.runtime_state.session_id, job_id=job_id) + return _tool_unexpected_error("shell_input", exc, job_id=job_id) async def _tool_interrupt( self, @@ -498,16 +382,14 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC output_path = (await self._require_resource().commands.tail(job_id)).output_path except (RuntimeError, ValueError) as exc: logger.warning( - "Failed to fetch output path for interrupted shell job %s in session %s: %s", + "Failed to fetch output path for interrupted shell job %s: %s", job_id, - self.runtime_state.session_id, exc, ) except Exception: logger.exception( - "Failed to fetch output path for interrupted shell job %s in session %s", + "Failed to fetch output path for interrupted shell job %s", job_id, - self.runtime_state.session_id, ) return _tagged_shell_observation( _metadata_dict( @@ -522,9 +404,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc, job_id=job_id) except Exception as exc: - return _tool_unexpected_error( - "shell_interrupt", exc, session_id=self.runtime_state.session_id, job_id=job_id - ) + return _tool_unexpected_error("shell_interrupt", exc, job_id=job_id) async def run_remote_script_complete( self, @@ -559,45 +439,6 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC inject_agent_stub_env=inject_agent_stub_env, ) - async def _allocate_workspace(self) -> tuple[str, str]: - for _attempt in range(_SESSION_ID_ATTEMPT_LIMIT): - session_id = _generate_session_id() - result = await self._run_internal_script_complete(_workspace_mkdir_script(session_id=session_id), cwd=None) - if result.exit_code == _WORKSPACE_COLLISION_EXIT_CODE: - continue - if result.exit_code != 0 or not result.output_complete: - raise RuntimeError( - f"Failed to create shell workspace {_workspace_cwd(session_id)}: " - + f"{result.status} exit_code={result.exit_code}" - ) - return session_id, _workspace_cwd(session_id) - raise RuntimeError("Failed to allocate a unique shell workspace session id after 256 attempts.") - - async def _bootstrap_workspace(self, session_id: str) -> None: - bootstrap_script = _workspace_bootstrap_script(self.config) - if not bootstrap_script: - return - workspace_cwd = _workspace_cwd_for_home(self._shell_home_dir(), session_id) - result = await self._run_internal_script_complete(bootstrap_script, cwd=workspace_cwd) - if result.exit_code != 0 or not result.output_complete: - raise RuntimeError( - f"Failed to bootstrap shell workspace {workspace_cwd}: {result.status} exit_code={result.exit_code}" - ) - - async def _cleanup_workspace_best_effort(self, session_id: str) -> None: - try: - _ = await self._run_internal_script_complete(_workspace_cleanup_script(session_id=session_id), cwd=None) - except (RuntimeError, ValueError) as exc: - logger.warning("Failed to remove shell workspace for session %s after create failure: %s", session_id, exc) - - async def _ensure_live_workspace_exists(self, session_id: str) -> None: - result = await self._run_internal_script_complete(_workspace_ensure_script(session_id=session_id), cwd=None) - if result.exit_code != 0 or not result.output_complete: - raise RuntimeError( - f"Failed to ensure shell workspace {_workspace_cwd_for_home(self._shell_home_dir(), session_id)} " - + f"exists: {result.status} exit_code={result.exit_code}" - ) - async def _run_internal_script_complete( self, script: str, @@ -613,36 +454,11 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC max_output_bytes=_REMOTE_COMPLETE_OUTPUT_MAX_BYTES, ) - def _require_resource(self) -> ShellResourceProtocol: - if self._shell_resource is None: - raise RuntimeError( - "DifyShellLayer requires an active shell resource inside resource_context(); " - + "enter the layer through Agenton or wrap direct hook/tool usage in resource_context()." - ) - return self._shell_resource + def _require_resource(self) -> RuntimeLease: + return self.deps.runtime.lease def _require_workspace_cwd(self) -> str: - session_id, _workspace_cwd = self._require_session_identity() - return _workspace_cwd_for_home(self._shell_home_dir(), session_id) - - def _require_session_identity(self) -> tuple[str, str]: - identity = self._try_session_identity() - if identity is None: - raise ValueError("DifyShellLayer runtime state is missing session_id or workspace_cwd.") - session_id, workspace_cwd = identity - expected_workspace = _workspace_cwd(session_id) - if workspace_cwd != expected_workspace: - raise ValueError( - f"DifyShellLayer runtime state has inconsistent workspace_cwd {workspace_cwd!r}; expected {expected_workspace!r}." - ) - return session_id, workspace_cwd - - def _try_session_identity(self) -> tuple[str, str] | None: - session_id = self.runtime_state.session_id - workspace_cwd = self.runtime_state.workspace_cwd - if session_id is None or workspace_cwd is None: - return None - return session_id, workspace_cwd + return self._require_resource().layout.workspace_dir def _ensure_tracked_job(self, job_id: str) -> None: if job_id not in self.runtime_state.job_ids: @@ -669,9 +485,8 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC await commands.delete(job_id, force=True) except RuntimeError as exc: logger.warning( - "Failed to delete shell job %s for session %s: %s", + "Failed to delete shell job %s: %s", job_id, - self.runtime_state.session_id, exc, ) @@ -679,23 +494,6 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC self.runtime_state.job_offsets = {} self.runtime_state.job_ids = [] - def _shell_home_dir(self) -> str: - return _shell_home_dir_for_agent_id( - self._require_current_execution_agent_id(), - shell_home_root=self.shell_home_root, - ) - - def _current_execution_agent_id(self) -> str | None: - execution_context_layer = self.deps.execution_context - execution_context = execution_context_layer.config if execution_context_layer is not None else None - return execution_context.agent_id if execution_context is not None else None - - def _require_current_execution_agent_id(self) -> str: - agent_id = self._current_execution_agent_id() - if agent_id is None: - raise ValueError("ShellLayer command execution requires execution_context.agent_id.") - return _validated_agent_home_segment(agent_id) - def _build_shell_command_env( self, *, @@ -703,7 +501,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC require_agent_stub_env: bool = False, ) -> dict[str, str]: env = _shell_config_env(self.config) - env["HOME"] = self._shell_home_dir() + env["HOME"] = self._require_resource().layout.home_dir if not include_agent_stub_env: return env execution_context_layer = self.deps.execution_context @@ -713,7 +511,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC agent_stub_drive_ref=self.config.agent_stub_drive_ref, execution_context=execution_context, token_factory=self.agent_stub_token_factory, - session_id=self.runtime_state.session_id, + session_id=None, ) if agent_stub_env is None: if not require_agent_stub_env: @@ -897,58 +695,19 @@ def _tool_unexpected_error( tool_name: str, exc: Exception, *, - session_id: str | None, job_id: str | None = None, ) -> ShellToolErrorObservation: # Unexpected Exception still becomes a tool observation so one shell tool # failure does not abort the agent loop, but it is logged with traceback for # debugging. BaseException is intentionally not caught by callers. logger.exception( - "Unexpected shell tool failure: tool=%s session_id=%s job_id=%s", + "Unexpected shell tool failure: tool=%s job_id=%s", tool_name, - session_id, job_id, ) return _tool_error_from_exception(exc, job_id=job_id) -def _generate_session_id() -> str: - time_component = int(time.time()) & _SESSION_TIME_HEX_MASK - random_component = secrets.token_hex(1) - if len(random_component) != _SESSION_RANDOM_HEX_LENGTH: - raise RuntimeError("Expected a one-byte random hex suffix for Dify shell session ids.") - return f"{time_component:05x}{random_component}" - - -def _workspace_cwd(session_id: str) -> str: - return f"{_WORKSPACE_ROOT}/{_validated_session_id(session_id)}" - - -def _normalize_shell_home_root(shell_home_root: str) -> str: - stripped = shell_home_root.strip().rstrip("/") - if not stripped: - raise ValueError("shell_home_root must not be empty") - if not stripped.startswith("/"): - raise ValueError("shell_home_root must be an absolute path") - return stripped - - -def _shell_home_dir_for_agent_id(agent_id: str | None, *, shell_home_root: str = "/home") -> str: - if agent_id is None: - raise ValueError("ShellLayer command execution requires execution_context.agent_id.") - return f"{_normalize_shell_home_root(shell_home_root)}/{_validated_agent_home_segment(agent_id)}" - - -def _validated_agent_home_segment(agent_id: str) -> str: - if agent_id in {".", ".."} or not _AGENT_HOME_SEGMENT_PATTERN.fullmatch(agent_id): - raise ValueError("execution_context.agent_id must be a safe single path segment for shell HOME.") - return agent_id - - -def _workspace_cwd_for_home(home_dir: str, session_id: str) -> str: - return f"{home_dir}/{_WORKSPACE_DIR_NAME}/{_validated_session_id(session_id)}" - - def _workspace_bootstrap_script(config: DifyShellLayerConfig) -> str: install_commands = [command for tool in config.cli_tools for command in tool.install_commands] if not install_commands: @@ -966,12 +725,6 @@ def _shell_config_export_lines(config: DifyShellLayerConfig) -> list[str]: for tool in config.cli_tools: for secret_ref in tool.secret_refs: lines.append(f'export {secret_ref.name}="${{{secret_ref.name}:-}}"') - if config.sandbox is not None: - if config.sandbox.provider: - lines.append(f"export DIFY_SANDBOX_PROVIDER={_shquote(config.sandbox.provider)}") - if config.sandbox.config: - sandbox_config = json.dumps(config.sandbox.config, ensure_ascii=True, sort_keys=True) - lines.append(f"export DIFY_SANDBOX_CONFIG_JSON={_shquote(sandbox_config)}") return lines @@ -992,35 +745,10 @@ def _wrap_user_script(script: str, config: DifyShellLayerConfig) -> str: return "\n".join([*lines, script]) -def _workspace_mkdir_script(*, session_id: str) -> str: - safe_session_id = _validated_session_id(session_id) - workspace_dir = f"$HOME/workspace/{safe_session_id}" - return ( - 'mkdir -p "$HOME/workspace"; ' - f'if mkdir "{workspace_dir}"; then exit 0; fi; ' - f'if [ -e "{workspace_dir}" ]; then exit {_WORKSPACE_COLLISION_EXIT_CODE}; fi; ' - "exit 1" - ) - - -def _workspace_cleanup_script(*, session_id: str) -> str: - return f'rm -rf -- "$HOME/workspace/{_validated_session_id(session_id)}"' - - -def _workspace_ensure_script(*, session_id: str) -> str: - return f'mkdir -p "$HOME/workspace/{_validated_session_id(session_id)}"' - - def _shquote(value: str) -> str: return "'" + value.replace("'", "'\\''") + "'" -def _validated_session_id(session_id: str) -> str: - if not _SESSION_ID_PATTERN.fullmatch(session_id): - raise ValueError("session_id must match the 5+2 lowercase hex format '<5 hex><2 hex>'.") - return session_id - - def _deduplicate_preserving_order(values: Sequence[str]) -> list[str]: seen: set[str] = set() result: list[str] = [] diff --git a/dify-agent/src/dify_agent/protocol/__init__.py b/dify-agent/src/dify_agent/protocol/__init__.py index 2fc5bd517ed..286dff9dafc 100644 --- a/dify-agent/src/dify_agent/protocol/__init__.py +++ b/dify-agent/src/dify_agent/protocol/__init__.py @@ -37,20 +37,26 @@ from .schemas import ( normalize_composition, utc_now, ) -from .sandbox import ( - RuntimeLayerSpec, - SandboxFileEntry, - SandboxListRequest, - SandboxListResponse, - SandboxLocator, - SandboxReadRequest, - SandboxReadResponse, - SandboxUploadRequest, - SandboxUploadResponse, - SandboxUploadedFile, - build_sandbox_locator_from_layer_specs, - build_sandbox_locator_from_run_request, - extract_runtime_layer_specs, +from .execution_binding import ( + CreateExecutionBindingRequest, + CreateExecutionBindingResponse, + DestroyExecutionBindingRequest, +) +from .home_snapshot import ( + CreateHomeSnapshotFromBindingRequest, + DeleteHomeSnapshotRequest, + HomeSnapshotResponse, + InitializeHomeSnapshotRequest, +) +from .workspace import ( + WorkspaceFileEntry, + WorkspaceListRequest, + WorkspaceListResponse, + WorkspaceReadRequest, + WorkspaceReadResponse, + WorkspaceUploadRequest, + WorkspaceUploadResponse, + WorkspaceUploadedFile, ) __all__ = [ @@ -60,13 +66,20 @@ __all__ = [ "CancelRunResponse", "CreateRunRequest", "CreateRunResponse", + "CreateExecutionBindingRequest", + "CreateExecutionBindingResponse", + "CreateHomeSnapshotFromBindingRequest", + "DeleteHomeSnapshotRequest", "DeferredToolCallPayload", "DeferredToolResultsPayload", "DIFY_AGENT_HISTORY_LAYER_ID", "DIFY_AGENT_MODEL_LAYER_ID", "DIFY_AGENT_OUTPUT_LAYER_ID", + "DestroyExecutionBindingRequest", "EmptyRunEventData", "LayerExitSignals", + "HomeSnapshotResponse", + "InitializeHomeSnapshotRequest", "PydanticAIStreamRunEvent", "RUN_EVENT_ADAPTER", "RunCancelledEvent", @@ -83,19 +96,14 @@ __all__ = [ "RunStatusResponse", "RunSucceededEvent", "RunSucceededEventData", - "RuntimeLayerSpec", - "SandboxFileEntry", - "SandboxListRequest", - "SandboxListResponse", - "SandboxLocator", - "SandboxReadRequest", - "SandboxReadResponse", - "SandboxUploadRequest", - "SandboxUploadResponse", - "SandboxUploadedFile", - "build_sandbox_locator_from_layer_specs", - "build_sandbox_locator_from_run_request", - "extract_runtime_layer_specs", + "WorkspaceFileEntry", + "WorkspaceListRequest", + "WorkspaceListResponse", + "WorkspaceReadRequest", + "WorkspaceReadResponse", + "WorkspaceUploadRequest", + "WorkspaceUploadResponse", + "WorkspaceUploadedFile", "normalize_composition", "utc_now", ] diff --git a/dify-agent/src/dify_agent/protocol/execution_binding.py b/dify-agent/src/dify_agent/protocol/execution_binding.py new file mode 100644 index 00000000000..29d212dd218 --- /dev/null +++ b/dify-agent/src/dify_agent/protocol/execution_binding.py @@ -0,0 +1,44 @@ +"""Private DTOs for persistent Execution Binding lifecycle operations.""" + +from typing import ClassVar + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class CreateExecutionBindingRequest(BaseModel): + tenant_id: str = Field(min_length=1) + agent_id: str = Field(min_length=1) + binding_id: str = Field(min_length=1) + workspace_id: str = Field(min_length=1) + existing_workspace_ref: str | None = None + home_snapshot_ref: str = Field(min_length=1) + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class CreateExecutionBindingResponse(BaseModel): + binding_ref: str = Field(min_length=1) + workspace_ref: str = Field(min_length=1) + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class DestroyExecutionBindingRequest(BaseModel): + binding_ref: str = Field(min_length=1) + destroy_workspace: bool + workspace_ref: str | None = None + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + @model_validator(mode="after") + def validate_workspace_ref(self) -> "DestroyExecutionBindingRequest": + if self.destroy_workspace and not self.workspace_ref: + raise ValueError("workspace_ref is required when destroy_workspace is true") + return self + + +__all__ = [ + "CreateExecutionBindingRequest", + "CreateExecutionBindingResponse", + "DestroyExecutionBindingRequest", +] diff --git a/dify-agent/src/dify_agent/protocol/home_snapshot.py b/dify-agent/src/dify_agent/protocol/home_snapshot.py new file mode 100644 index 00000000000..0b4ddad9adc --- /dev/null +++ b/dify-agent/src/dify_agent/protocol/home_snapshot.py @@ -0,0 +1,42 @@ +"""Private DTOs for immutable Home Snapshot operations.""" + +from typing import ClassVar + +from pydantic import BaseModel, ConfigDict, Field + + +class InitializeHomeSnapshotRequest(BaseModel): + tenant_id: str = Field(min_length=1) + agent_id: str = Field(min_length=1) + home_snapshot_id: str = Field(min_length=1) + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class CreateHomeSnapshotFromBindingRequest(BaseModel): + tenant_id: str = Field(min_length=1) + agent_id: str = Field(min_length=1) + home_snapshot_id: str = Field(min_length=1) + backend_binding_ref: str = Field(min_length=1) + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class DeleteHomeSnapshotRequest(BaseModel): + snapshot_ref: str = Field(min_length=1) + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class HomeSnapshotResponse(BaseModel): + snapshot_ref: str = Field(min_length=1) + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +__all__ = [ + "CreateHomeSnapshotFromBindingRequest", + "DeleteHomeSnapshotRequest", + "HomeSnapshotResponse", + "InitializeHomeSnapshotRequest", +] diff --git a/dify-agent/src/dify_agent/protocol/sandbox.py b/dify-agent/src/dify_agent/protocol/sandbox.py deleted file mode 100644 index db375e6a7b5..00000000000 --- a/dify-agent/src/dify_agent/protocol/sandbox.py +++ /dev/null @@ -1,254 +0,0 @@ -"""Public sandbox DTOs shared by the API and Dify Agent backends. - -The sandbox file APIs must rebuild only the minimum runtime needed to re-enter a -prior shell session: ``dify.execution_context`` for Dify-owned identity and -``dify.shell`` for the sandbox workspace itself. ``SandboxLocator`` therefore -contains a safe composition subset plus the matching filtered session snapshot. -Credential-bearing or runtime-only tool layers are intentionally excluded from -persisted runtime specs and from sandbox locators. -""" - -from __future__ import annotations - -from typing import ClassVar, Literal, cast - -from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot -from dify_agent.layers.dify_core_tools import DIFY_CORE_TOOLS_LAYER_TYPE_ID -from dify_agent.layers.dify_plugin import DIFY_PLUGIN_LLM_LAYER_TYPE_ID, DIFY_PLUGIN_TOOLS_LAYER_TYPE_ID -from pydantic import BaseModel, ConfigDict, Field, JsonValue - -from .schemas import CreateRunRequest, RunComposition, RunLayerSpec - -_SENSITIVE_LAYER_TYPES = frozenset( - { - DIFY_PLUGIN_LLM_LAYER_TYPE_ID, - DIFY_PLUGIN_TOOLS_LAYER_TYPE_ID, - DIFY_CORE_TOOLS_LAYER_TYPE_ID, - } -) - - -class RuntimeLayerSpec(BaseModel): - """Persistable non-sensitive layer spec derived from a run composition. - - API-side runtime-session rows store these specs so later cleanup or sandbox - requests can rebuild the minimal layer graph without persisting model or - tool credentials. - """ - - name: str - type: str - deps: dict[str, str] = Field(default_factory=dict) - metadata: dict[str, JsonValue] = Field(default_factory=dict) - config: JsonValue = None - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxLocator(BaseModel): - """Safe subset of one prior run request needed to re-enter a sandbox shell.""" - - composition: RunComposition - session_snapshot: CompositorSessionSnapshot - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxFileEntry(BaseModel): - """One directory entry returned by ``/sandbox/files/list``.""" - - name: str - type: Literal["file", "dir", "symlink", "other"] - size: int | None = None - mtime: int | None = None - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxListRequest(BaseModel): - """Request body for listing a sandbox directory.""" - - locator: SandboxLocator - path: str = "." - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxListResponse(BaseModel): - """Structured sandbox directory listing.""" - - path: str - entries: list[SandboxFileEntry] - truncated: bool - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxReadRequest(BaseModel): - """Request body for reading a sandbox file preview.""" - - locator: SandboxLocator - path: str - max_bytes: int = 262144 - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxReadResponse(BaseModel): - """Text preview returned by ``/sandbox/files/read``.""" - - path: str - size: int | None = None - truncated: bool - binary: bool - text: str | None = None - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxUploadedFile(BaseModel): - """Canonical ToolFile mapping returned after sandbox upload.""" - - transfer_method: Literal["tool_file"] = "tool_file" - reference: str - download_url: str - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxUploadRequest(BaseModel): - """Request body for uploading one sandbox file through the Agent Stub.""" - - locator: SandboxLocator - path: str - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -class SandboxUploadResponse(BaseModel): - """Result returned after sandbox upload creates a ToolFile mapping.""" - - path: str - file: SandboxUploadedFile - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -def extract_runtime_layer_specs(composition: RunComposition) -> list[RuntimeLayerSpec]: - """Project a run composition into the persistable non-sensitive layer list.""" - specs: list[RuntimeLayerSpec] = [] - for layer in composition.layers: - if layer.type in _SENSITIVE_LAYER_TYPES: - continue - config_value: JsonValue = None - if isinstance(layer.config, BaseModel): - config_value = layer.config.model_dump(mode="json", warnings=False) - else: - config_value = cast(JsonValue, layer.config) - specs.append( - RuntimeLayerSpec( - name=layer.name, - type=layer.type, - deps=dict(layer.deps), - metadata=dict(layer.metadata), - config=config_value, - ) - ) - return specs - - -def build_sandbox_locator_from_run_request(request: CreateRunRequest) -> SandboxLocator: - """Build a safe sandbox locator from a full create-run request. - - Raises: - ValueError: if the request has no resumable session snapshot or lacks the - execution-context/shell layers needed for sandbox access. - """ - if request.session_snapshot is None: - raise ValueError("Sandbox locator requires a non-empty session_snapshot.") - return build_sandbox_locator_from_layer_specs( - layer_specs=extract_runtime_layer_specs(request.composition), - session_snapshot=request.session_snapshot, - ) - - -def build_sandbox_locator_from_layer_specs( - *, - layer_specs: list[RuntimeLayerSpec], - session_snapshot: CompositorSessionSnapshot, -) -> SandboxLocator: - """Build a sandbox locator from persisted runtime specs plus a saved snapshot.""" - if not layer_specs: - raise ValueError("Sandbox locator requires persisted runtime layer specs.") - - for spec in layer_specs: - if spec.type in _SENSITIVE_LAYER_TYPES: - raise ValueError(f"Sandbox locator runtime specs must not include sensitive layer type {spec.type!r}.") - - execution_context_index = next( - (index for index, spec in enumerate(layer_specs) if spec.name == "execution_context"), - None, - ) - shell_index = next((index for index, spec in enumerate(layer_specs) if spec.name == "shell"), None) - if execution_context_index is None: - raise ValueError("Sandbox locator requires an 'execution_context' runtime layer spec.") - if shell_index is None: - raise ValueError("Sandbox locator requires a 'shell' runtime layer spec.") - if execution_context_index > shell_index: - raise ValueError("Sandbox locator requires 'execution_context' to appear before 'shell'.") - - execution_context_spec = layer_specs[execution_context_index] - shell_spec = layer_specs[shell_index] - if shell_spec.deps.get("execution_context") != execution_context_spec.name: - raise ValueError("Sandbox shell layer must depend on the execution_context layer.") - - kept_specs = [execution_context_spec, shell_spec] - kept_names = [spec.name for spec in kept_specs] - snapshot_layers = [layer for layer in session_snapshot.layers if layer.name in set(kept_names)] - if [layer.name for layer in snapshot_layers] != kept_names: - raise ValueError("Sandbox locator session_snapshot must contain execution_context and shell layers in order.") - - return SandboxLocator( - composition=RunComposition( - schema_version=1, - layers=[ - RunLayerSpec( - name=spec.name, - type=spec.type, - deps=dict(spec.deps), - metadata=dict(spec.metadata), - config=spec.config, - ) - for spec in kept_specs - ], - ), - session_snapshot=CompositorSessionSnapshot( - schema_version=session_snapshot.schema_version, - layers=[ - LayerSessionSnapshot( - name=layer.name, - lifecycle_state=layer.lifecycle_state, - runtime_state=dict(layer.runtime_state), - ) - for layer in snapshot_layers - ], - ), - ) - - -__all__ = [ - "RuntimeLayerSpec", - "SandboxFileEntry", - "SandboxListRequest", - "SandboxListResponse", - "SandboxLocator", - "SandboxReadRequest", - "SandboxReadResponse", - "SandboxUploadRequest", - "SandboxUploadResponse", - "SandboxUploadedFile", - "build_sandbox_locator_from_layer_specs", - "build_sandbox_locator_from_run_request", - "extract_runtime_layer_specs", -] diff --git a/dify-agent/src/dify_agent/protocol/schemas.py b/dify-agent/src/dify_agent/protocol/schemas.py index b4ea3843918..ae4c9324294 100644 --- a/dify-agent/src/dify_agent/protocol/schemas.py +++ b/dify-agent/src/dify_agent/protocol/schemas.py @@ -23,13 +23,9 @@ by ``DIFY_AGENT_OUTPUT_LAYER_ID``. Request-level ``on_exit`` signals decide whether each active layer is suspended or deleted when the run exits, with suspend as the default so successful terminal events can include resumable snapshots. Successful runs always publish the resumable Agenton session snapshot -on the terminal ``run_succeeded`` event together with exactly one of the final -JSON-safe ``output`` or a deferred external ``deferred_tool_call`` payload. A -lifecycle-only run may also succeed with ``output = null`` and ``usage = null`` -when the composition intentionally omits the reserved model layer and only -replays layer enter/exit work from a supplied snapshot. That lets consumers -treat terminal success events as complete run summaries without a separate pause -protocol. Session snapshots carry only layer lifecycle/runtime state in +on the terminal ``run_succeeded`` event together with either the final JSON-safe +``output`` or a deferred external ``deferred_tool_call`` payload. Session +snapshots carry only layer lifecycle/runtime state in compositor order; they do not persist output-layer config. Resumed structured-output runs therefore must resubmit the same ``output`` layer in ``composition.layers[]`` so snapshot layer name/order still matches the diff --git a/dify-agent/src/dify_agent/protocol/workspace.py b/dify-agent/src/dify_agent/protocol/workspace.py new file mode 100644 index 00000000000..ad1d909fdee --- /dev/null +++ b/dify-agent/src/dify_agent/protocol/workspace.py @@ -0,0 +1,84 @@ +"""Private Workspace file DTOs resolved through an Execution Binding ref.""" + +from typing import ClassVar, Literal + +from pydantic import BaseModel, ConfigDict, Field + +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig + + +class WorkspaceFileEntry(BaseModel): + name: str + type: Literal["file", "dir", "symlink", "other"] + size: int | None = None + mtime: int | None = None + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceListRequest(BaseModel): + backend_binding_ref: str = Field(min_length=1) + path: str = "." + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceListResponse(BaseModel): + path: str + entries: list[WorkspaceFileEntry] + truncated: bool + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceReadRequest(BaseModel): + backend_binding_ref: str = Field(min_length=1) + path: str + max_bytes: int = 262144 + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceReadResponse(BaseModel): + path: str + size: int | None = None + truncated: bool + binary: bool + text: str | None = None + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceUploadedFile(BaseModel): + transfer_method: Literal["tool_file"] = "tool_file" + reference: str + download_url: str + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceUploadRequest(BaseModel): + backend_binding_ref: str = Field(min_length=1) + path: str + execution_context: DifyExecutionContextLayerConfig + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceUploadResponse(BaseModel): + path: str + file: WorkspaceUploadedFile + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +__all__ = [ + "WorkspaceFileEntry", + "WorkspaceListRequest", + "WorkspaceListResponse", + "WorkspaceReadRequest", + "WorkspaceReadResponse", + "WorkspaceUploadRequest", + "WorkspaceUploadResponse", + "WorkspaceUploadedFile", +] diff --git a/dify-agent/src/dify_agent/runtime/compositor_factory.py b/dify-agent/src/dify_agent/runtime/compositor_factory.py index b81cf737478..429c7fce57a 100644 --- a/dify-agent/src/dify_agent/runtime/compositor_factory.py +++ b/dify-agent/src/dify_agent/runtime/compositor_factory.py @@ -3,12 +3,13 @@ Only explicitly allowed provider type ids are constructible here. The default provider set contains prompt layers, the optional pydantic-ai history layer, the state-free Dify structured output layer, the optional Dify ask-human layer, the -Dify execution-context layer, the stateful Dify shell layer, and the Dify -plugin/knowledge business-layer family: +Dify execution-context and runtime-resource layers, the stateful Dify shell +layer, and the Dify plugin/knowledge business-layer family: - ``dify.config`` for Agent Soul-backed config assets + eager pull, - ``dify.execution_context`` for shared tenant/user/run daemon context, -- ``dify.shell`` for shellctl-backed shell job control, +- ``dify.runtime`` for operation-scoped RuntimeLease acquisition, +- ``dify.shell`` for command/file capabilities from the active RuntimeLease, - ``dify.plugin.llm`` for plugin-backed model selection, - ``dify.plugin.tools`` for prepared plugin tool exposure, and - ``dify.core.tools`` for API-routed Dify tool exposure, and @@ -16,10 +17,9 @@ plugin/knowledge business-layer family: Public DTOs provide Dify context plus plugin/model/tool data, while server-only plugin daemon settings and Dify API inner settings are injected through provider -factories. An already-selected shell provider (shellctl or enterprise, built at -the runtime boundary) and the Agent Stub URL/token factory are injected for -``DifyShellLayer``; a ``None`` shell provider leaves the shell layer disabled. -The resulting ``Compositor`` +factories. The deployment-selected runtime backend profile is injected only +into the Runtime layer. Shell receives Agent Stub settings and consumes the +active lease's data plane. The resulting ``Compositor`` remains Agenton state-only at the snapshot boundary: live resources such as HTTP clients are injected by runtime-owned providers, may be held on active layer instances inside ``resource_context()``, and never enter session @@ -29,10 +29,7 @@ snapshots. from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, cast - -if TYPE_CHECKING: - from dify_agent.adapters.shell.protocols import ShellProviderProtocol +from typing import Any, cast from pydantic_ai.messages import UserContent @@ -55,8 +52,11 @@ from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer from dify_agent.layers.knowledge.configs import DifyKnowledgeBaseLayerConfig from dify_agent.layers.knowledge.layer import DifyKnowledgeBaseLayer from dify_agent.layers.output.output_layer import DifyOutputLayer +from dify_agent.layers.runtime.configs import DifyRuntimeLayerConfig +from dify_agent.layers.runtime.layer import DifyRuntimeLayer from dify_agent.layers.shell.configs import DifyShellLayerConfig from dify_agent.layers.shell.layer import DifyShellLayer +from dify_agent.runtime_backend import RuntimeBackendProfile type DifyAgentLayerProvider = LayerProvider[Any] @@ -67,14 +67,13 @@ def create_default_layer_providers( plugin_daemon_api_key: str = "", inner_api_url: str = "http://localhost:5001", inner_api_key: str = "", - shell_provider: ShellProviderProtocol | None = None, - shell_home_root: str = "/home", + runtime_backend_profile: RuntimeBackendProfile | None = None, shell_redact_patterns: list[str] | None = None, agent_stub_api_base_url: str | None = None, agent_stub_token_factory: ShellAgentStubTokenFactory | None = None, ) -> tuple[DifyAgentLayerProvider, ...]: """Return the server provider set of safe config-constructible layers.""" - return ( + providers: list[DifyAgentLayerProvider] = [ LayerProvider.from_layer_type(PromptLayer), LayerProvider.from_layer_type(PydanticAIHistoryLayer), LayerProvider.from_layer_type(DifyOutputLayer), @@ -93,8 +92,6 @@ def create_default_layer_providers( layer_type=DifyShellLayer, create=lambda config: DifyShellLayer.from_config_with_settings( DifyShellLayerConfig.model_validate(config), - shell_provider=shell_provider, - shell_home_root=shell_home_root, shell_redact_patterns=shell_redact_patterns or [], agent_stub_api_base_url=agent_stub_api_base_url, agent_stub_token_factory=agent_stub_token_factory, @@ -125,7 +122,20 @@ def create_default_layer_providers( inner_api_key=inner_api_key, ), ), - ) + ] + if runtime_backend_profile is not None: + providers.extend( + [ + LayerProvider.from_factory( + layer_type=DifyRuntimeLayer, + create=lambda config: DifyRuntimeLayer.from_config_with_backend( + DifyRuntimeLayerConfig.model_validate(config), + backend=runtime_backend_profile.execution_bindings, + ), + ), + ] + ) + return tuple(providers) def build_pydantic_ai_compositor( diff --git a/dify-agent/src/dify_agent/runtime/runner.py b/dify-agent/src/dify_agent/runtime/runner.py index 8fec8d5fbf8..71868a2324c 100644 --- a/dify-agent/src/dify_agent/runtime/runner.py +++ b/dify-agent/src/dify_agent/runtime/runner.py @@ -1,18 +1,14 @@ """Runtime execution for one scheduled Dify Agent run. The runner is storage-agnostic: it normalizes the public Dify composition into -Agenton's graph/config split and chooses one of two execution modes after the -composition is normalized and the ``on_exit`` policy is validated: +Agenton's graph/config split and executes one model run after the ``on_exit`` +policy is validated: - model runs: enter a fresh ``CompositorRun`` (or resume one from a snapshot), render the current Dify system prompts into temporary ``message_history``, run pydantic-ai with either the current ``run.user_prompts`` or deferred external tool results, emit raw stream events with agent-message delta annotations, apply request-level ``on_exit`` signals, and publish a terminal success or failure event; -- lifecycle-only runs: enter from a supplied snapshot, apply request-level - ``on_exit`` signals, exit without invoking a model, and succeed with explicit - ``output = null`` and ``usage = null``. - The Pydantic AI model is resolved from the active Agenton layer named by ``DIFY_AGENT_MODEL_LAYER_ID``. An optional history layer contributes stored message history only through session state; successful model runs append only @@ -63,7 +59,7 @@ from dify_agent.protocol.schemas import ( from dify_agent.runtime.agent_factory import create_agent, normalize_user_input from dify_agent.runtime.agenton_validation import is_agenton_enter_validation_runtime_error from dify_agent.runtime.compositor_factory import build_pydantic_ai_compositor, create_default_layer_providers -from dify_agent.adapters.shell.protocols import SandboxExpiredError +from dify_agent.runtime_backend import BindingLostError from dify_agent.runtime.event_sink import ( RunEventSink, emit_pydantic_ai_event, @@ -120,8 +116,8 @@ def _run_failed_error_payload(exc: Exception) -> tuple[str, str | None]: message = str(exc) or type(exc).__name__ reason: str | None = None - if isinstance(exc, SandboxExpiredError): - return message, "sandbox_expired" + if isinstance(exc, BindingLostError): + return message, "binding_lost" if isinstance(exc, ModelHTTPError): body = exc.body @@ -234,7 +230,7 @@ class AgentRunRunner: await self.sink.update_status(self.run_id, "succeeded") async def _run_agent(self) -> RunSuccessOutcome: - """Run the normalized request in model or lifecycle-only mode. + """Run the normalized request through the model path. Known request-shaped Agenton enter-time failures are normalized to ``AgentRunValidationError``. That includes the existing small class of @@ -262,49 +258,9 @@ class AgentRunRunner: raise AgentRunValidationError(str(exc)) from exc if not _has_model_layer(self.request): - return await self._run_lifecycle_only(compositor=compositor, layer_configs=layer_configs) + raise AgentRunValidationError(f"Missing required '{DIFY_AGENT_MODEL_LAYER_ID}' layer.") return await self._run_model(compositor=compositor, layer_configs=layer_configs) - async def _run_lifecycle_only( - self, - *, - compositor: Any, - layer_configs: dict[str, LayerConfigInput], - ) -> RunSuccessOutcome: - """Replay only layer lifecycle work for a no-LLM composition plus snapshot.""" - if self.request.session_snapshot is None: - raise AgentRunValidationError( - f"Missing '{DIFY_AGENT_MODEL_LAYER_ID}' requires a session_snapshot for lifecycle-only runs." - ) - if self.request.deferred_tool_results is not None: - raise AgentRunValidationError( - f"Deferred tool results require the reserved '{DIFY_AGENT_MODEL_LAYER_ID}' layer." - ) - - entered_run = False - try: - async with compositor.enter(configs=layer_configs, session_snapshot=self.request.session_snapshot) as run: - entered_run = True - apply_layer_exit_signals(run, self.request.on_exit) - except RuntimeError as exc: - if not entered_run and is_agenton_enter_validation_runtime_error(exc): - raise AgentRunValidationError(str(exc)) from exc - raise - except ValueError as exc: - if not entered_run: - raise AgentRunValidationError(str(exc)) from exc - raise - - if run.session_snapshot is None: - raise RuntimeError("Agenton run did not produce a session snapshot after exit.") - return RunSuccessOutcome( - result_kind="output", - output=None, - deferred_tool_call=None, - session_snapshot=run.session_snapshot, - usage=None, - ) - async def _run_model( self, *, diff --git a/dify-agent/src/dify_agent/runtime_backend/__init__.py b/dify-agent/src/dify_agent/runtime_backend/__init__.py new file mode 100644 index 00000000000..e39d736e084 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/__init__.py @@ -0,0 +1,63 @@ +"""Public contracts for deployment-selected runtime backends.""" + +from .errors import ( + BindingAcquireError, + BindingCreateError, + BindingDestroyError, + BindingLostError, + HomeSnapshotCreateError, + HomeSnapshotNotFoundError, + RuntimeBackendError, + SharedWorkspaceUnsupportedError, + WorkspaceFileTooLargeError, + WorkspacePathError, + WorkspacePreservationUnsupportedError, + WorkspaceUnavailableError, +) +from .protocols import ( + ExecutionBindingAllocation, + ExecutionBindingBackend, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + FileSystem, + HomeSnapshotBackend, + HomeSnapshotCreateSpec, + InitializeHomeSnapshotSpec, + RuntimeBackendProfile, + RuntimeLayout, + RuntimeLease, + WorkspaceFileContent, + WorkspaceFileEntry, + WorkspaceListResult, + WorkspaceReadResult, +) + +__all__ = [ + "BindingAcquireError", + "BindingCreateError", + "BindingDestroyError", + "BindingLostError", + "ExecutionBindingAllocation", + "ExecutionBindingBackend", + "ExecutionBindingCreateSpec", + "ExecutionBindingDestroySpec", + "FileSystem", + "HomeSnapshotBackend", + "HomeSnapshotCreateError", + "HomeSnapshotCreateSpec", + "HomeSnapshotNotFoundError", + "InitializeHomeSnapshotSpec", + "RuntimeBackendError", + "RuntimeBackendProfile", + "RuntimeLayout", + "RuntimeLease", + "SharedWorkspaceUnsupportedError", + "WorkspaceFileContent", + "WorkspaceFileEntry", + "WorkspaceFileTooLargeError", + "WorkspaceListResult", + "WorkspacePathError", + "WorkspacePreservationUnsupportedError", + "WorkspaceReadResult", + "WorkspaceUnavailableError", +] diff --git a/dify-agent/src/dify_agent/runtime_backend/e2b.py b/dify-agent/src/dify_agent/runtime_backend/e2b.py new file mode 100644 index 00000000000..17414249f0d --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/e2b.py @@ -0,0 +1,405 @@ +"""E2B backend adapters with shellctl as the command and file data plane. + +Dify API persists only opaque Home Snapshot, Binding, and Workspace backend +refs. This adapter maps those refs to E2B resources internally. API keys, +traffic tokens, SDK objects, shellctl clients, and ``RuntimeLease`` objects stay +operation-local and are never serialized into Agenton state. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Literal, Protocol, cast + +import httpx2 as httpx + +from dify_agent.adapters.shell.protocols import ShellCommandProtocol +from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol +from dify_agent.runtime_backend.errors import ( + BindingAcquireError, + BindingCreateError, + BindingDestroyError, + BindingLostError, + HomeSnapshotCreateError, + SharedWorkspaceUnsupportedError, + WorkspacePreservationUnsupportedError, +) +from dify_agent.runtime_backend.protocols import ( + ExecutionBindingAllocation, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + FileSystem, + HomeSnapshotCreateSpec, + InitializeHomeSnapshotSpec, + RuntimeLayout, + RuntimeLease, +) +from dify_agent.runtime_backend.shellctl import ShellctlRuntimeLease, create_owned_shellctl_lease + +if TYPE_CHECKING: + from e2b.connection_config import ApiParams + +E2B_MAX_ACTIVE_TIMEOUT_SECONDS = 60 * 60 + + +class _E2BControlPlaneNotFoundError(RuntimeError): + """Typed boundary error for SDK resources that no longer exist.""" + + +class _E2BFileSystem(Protocol): + async def make_dir(self, path: str) -> bool: ... + + async def exists(self, path: str) -> bool: ... + + async def remove(self, path: str) -> None: ... + + +class _E2BSnapshotInfo(Protocol): + snapshot_id: str + names: list[str] + + +class _E2BSandbox(Protocol): + sandbox_id: str + traffic_access_token: str | None + files: _E2BFileSystem + + def get_host(self, port: int) -> str: ... + + async def pause(self, keep_memory: bool = True) -> bool: ... + + async def kill(self) -> bool: ... + + async def create_snapshot(self, name: str | None = None) -> _E2BSnapshotInfo: ... + + +class E2BControlPlane(Protocol): + async def create( + self, + template: str, + *, + timeout: int, + metadata: dict[str, str], + on_timeout: Literal["kill", "pause"], + ) -> _E2BSandbox: ... + + async def connect(self, handle: str, *, timeout: int) -> _E2BSandbox: ... + + async def kill(self, handle: str) -> bool: ... + + async def delete_snapshot(self, snapshot_ref: str) -> bool: ... + + +@dataclass(frozen=True, slots=True) +class E2BSDKControlPlane: + """Stateless async E2B SDK boundary configured with one deployment API key. + + SDK Sandbox objects are returned only to operation-local backend adapters. + Native not-found exceptions are normalized so adapters can distinguish + confirmed resource loss from transient acquisition and cleanup failures. + """ + + api_key: str + + def _options(self) -> ApiParams: + return {"api_key": self.api_key} + + async def create( + self, + template: str, + *, + timeout: int, + metadata: dict[str, str], + on_timeout: Literal["kill", "pause"], + ) -> _E2BSandbox: + from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException + + try: + return cast( + _E2BSandbox, + cast( + object, + await AsyncSandbox.create( + template, + timeout=timeout, + metadata=metadata, + lifecycle={"on_timeout": on_timeout, "auto_resume": False}, + **self._options(), + ), + ), + ) + except (SandboxNotFoundException, NotFoundException) as exc: + raise _E2BControlPlaneNotFoundError(str(exc)) from exc + + async def connect(self, handle: str, *, timeout: int) -> _E2BSandbox: + from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException + + try: + return cast( + _E2BSandbox, + cast(object, await AsyncSandbox.connect(handle, timeout=timeout, **self._options())), + ) + except (SandboxNotFoundException, NotFoundException) as exc: + raise _E2BControlPlaneNotFoundError(str(exc)) from exc + + async def kill(self, handle: str) -> bool: + from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException + + try: + return await AsyncSandbox.kill(handle, **self._options()) + except (SandboxNotFoundException, NotFoundException) as exc: + raise _E2BControlPlaneNotFoundError(str(exc)) from exc + + async def delete_snapshot(self, snapshot_ref: str) -> bool: + from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException + + try: + return await AsyncSandbox.delete_snapshot(snapshot_ref, **self._options()) + except (SandboxNotFoundException, NotFoundException) as exc: + raise _E2BControlPlaneNotFoundError(str(exc)) from exc + + +@dataclass(slots=True) +class E2BHomeSnapshotBackend: + """Implement immutable Home Snapshot operations with E2B snapshots. + + Initialization snapshots the prepared deployment template and releases its + temporary E2B resource. Build Apply snapshots the E2B resource behind the + supplied ``RuntimeLease``. Dify API stores the returned value as an opaque + backend ref; this adapter keeps no cross-request state. + """ + + control_plane: E2BControlPlane + template: str + active_timeout_seconds: int + home_dir: str = "/home/dify" + + async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str: + sandbox: _E2BSandbox | None = None + try: + sandbox = await self.control_plane.create( + self.template, + timeout=self.active_timeout_seconds, + metadata={ + "dify.resource": "home-snapshot-initialize", + "dify.tenant_id": spec.tenant_id, + "dify.agent_id": spec.agent_id, + "dify.home_snapshot_id": spec.home_snapshot_id, + }, + on_timeout="kill", + ) + _ = await sandbox.files.make_dir(self.home_dir) + snapshot = await sandbox.create_snapshot() + return snapshot.snapshot_id + except BaseException as exc: + if isinstance(exc, Exception): + raise HomeSnapshotCreateError(str(exc)) from exc + raise + finally: + if sandbox is not None: + try: + _ = await sandbox.kill() + except BaseException: + pass + + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + """Create an immutable E2B snapshot from the source Binding's active lease.""" + del spec + if not isinstance(source, E2BRuntimeLease): + raise HomeSnapshotCreateError("E2B Home Snapshot requires an E2B RuntimeLease") + try: + snapshot = await source.sandbox.create_snapshot() + return snapshot.snapshot_id + except BaseException as exc: + if isinstance(exc, Exception): + raise HomeSnapshotCreateError(str(exc)) from exc + raise + + async def delete(self, snapshot_ref: str) -> None: + try: + _ = await self.control_plane.delete_snapshot(snapshot_ref) + except _E2BControlPlaneNotFoundError: + return + except Exception as exc: + raise BindingDestroyError(str(exc)) from exc + + +@dataclass(slots=True) +class E2BExecutionBindingBackend: + """Implement Execution Binding operations with E2B and shellctl. + + In this backend one physical E2B resource represents both a Binding and its + Workspace, so their opaque refs have the same value. Materialized Home and + Workspace remain distinct logical resources even though E2B couples their + physical lifecycle. Active timeout pauses the resource and is not a + resource-age TTL. + """ + + control_plane: E2BControlPlane + active_timeout_seconds: int + shellctl_auth_token: str = "" + shellctl_port: int = 5004 + layout: RuntimeLayout = field( + default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace") + ) + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + """Create one paused E2B resource from an immutable Home Snapshot ref.""" + if spec.existing_workspace_ref is not None: + raise SharedWorkspaceUnsupportedError("current E2B backend cannot attach to an existing Workspace") + sandbox: _E2BSandbox | None = None + try: + sandbox = await self.control_plane.create( + spec.home_snapshot_ref, + timeout=self.active_timeout_seconds, + metadata={ + "dify.resource": "runtime-sandbox", + "dify.binding_id": spec.binding_id, + "dify.workspace_id": spec.workspace_id, + "dify.tenant_id": spec.tenant_id, + "dify.agent_id": spec.agent_id, + }, + on_timeout="pause", + ) + if await sandbox.files.exists(self.layout.workspace_dir): + await sandbox.files.remove(self.layout.workspace_dir) + _ = await sandbox.files.make_dir(self.layout.workspace_dir) + sandbox_id = sandbox.sandbox_id + _ = await sandbox.pause(keep_memory=True) + return ExecutionBindingAllocation(binding_ref=sandbox_id, workspace_ref=sandbox_id) + except BaseException as exc: + if sandbox is not None: + try: + _ = await sandbox.kill() + except BaseException: + pass + if isinstance(exc, Exception): + if isinstance(exc, (BindingCreateError, SharedWorkspaceUnsupportedError)): + raise + raise BindingCreateError(str(exc)) from exc + raise + + async def acquire(self, binding_ref: str) -> RuntimeLease: + """Acquire operation-scoped shellctl access for an opaque Binding ref.""" + sandbox: _E2BSandbox | None = None + try: + sandbox = await self.control_plane.connect(binding_ref, timeout=self.active_timeout_seconds) + if not await sandbox.files.exists(self.layout.workspace_dir): + raise BindingLostError(f"E2B Binding {binding_ref!r} no longer contains its Workspace") + return await self._lease(sandbox) + except _E2BControlPlaneNotFoundError as exc: + raise BindingLostError(f"E2B Binding {binding_ref!r} no longer exists") from exc + except BindingLostError: + await _best_effort_pause(sandbox) + raise + except BaseException as exc: + await _best_effort_pause(sandbox) + if isinstance(exc, Exception): + raise BindingAcquireError(str(exc)) from exc + raise + + async def release(self, lease: RuntimeLease) -> None: + """Close operation-local transports and pause the physical E2B resource.""" + if not isinstance(lease, E2BRuntimeLease): + raise TypeError("E2BExecutionBindingBackend can only release its own RuntimeLease") + close_error: Exception | None = None + try: + await lease.data_plane.close() + except Exception as exc: + close_error = exc + try: + _ = await lease.sandbox.pause(keep_memory=True) + except Exception as exc: + raise BindingAcquireError(str(exc)) from exc + if close_error is not None: + raise BindingAcquireError(str(close_error)) from close_error + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + """Destroy the coupled physical Binding and Workspace idempotently.""" + if not spec.destroy_workspace: + raise WorkspacePreservationUnsupportedError( + "current E2B backend cannot destroy a Binding while preserving its Workspace" + ) + if spec.workspace_ref != spec.binding_ref: + raise BindingDestroyError("E2B Workspace ref must equal its Binding ref") + try: + _ = await self.control_plane.kill(spec.binding_ref) + except _E2BControlPlaneNotFoundError: + return + except Exception as exc: + raise BindingDestroyError(str(exc)) from exc + + async def _lease(self, sandbox: _E2BSandbox) -> "E2BRuntimeLease": + entrypoint = f"https://{sandbox.get_host(self.shellctl_port)}" + traffic_token = sandbox.traffic_access_token + headers = {"X-Access-Token": traffic_token} if isinstance(traffic_token, str) and traffic_token else {} + http_client = httpx.AsyncClient( + base_url=entrypoint, + headers=headers, + follow_redirects=True, + timeout=httpx.Timeout(60.0), + ) + + def client_factory() -> ShellctlClientProtocol: + from shellctl.client import ShellctlClient + + return cast( + ShellctlClientProtocol, + cast( + object, + ShellctlClient(entrypoint, token=self.shellctl_auth_token, client=http_client), + ), + ) + + data_plane = await create_owned_shellctl_lease( + handle=sandbox.sandbox_id, + layout=self.layout, + entrypoint=entrypoint, + token=self.shellctl_auth_token, + client_factory=client_factory, + owned_transport=http_client, + ) + return E2BRuntimeLease(sandbox=sandbox, data_plane=data_plane) + + +@dataclass(slots=True) +class E2BRuntimeLease: + """Invocation-local E2B SDK object plus the owned shellctl data-plane lease.""" + + sandbox: _E2BSandbox + data_plane: ShellctlRuntimeLease + + @property + def handle(self) -> str: + return self.data_plane.handle + + @property + def layout(self) -> RuntimeLayout: + return self.data_plane.layout + + @property + def commands(self) -> ShellCommandProtocol: + return self.data_plane.commands + + @property + def files(self) -> FileSystem: + return self.data_plane.files + + +async def _best_effort_pause(sandbox: _E2BSandbox | None) -> None: + if sandbox is None: + return + try: + _ = await sandbox.pause(keep_memory=True) + except BaseException: + pass + + +__all__ = [ + "E2B_MAX_ACTIVE_TIMEOUT_SECONDS", + "E2BControlPlane", + "E2BExecutionBindingBackend", + "E2BHomeSnapshotBackend", + "E2BSDKControlPlane", + "E2BRuntimeLease", +] diff --git a/dify-agent/src/dify_agent/runtime_backend/enterprise.py b/dify-agent/src/dify_agent/runtime_backend/enterprise.py new file mode 100644 index 00000000000..a7de33d2514 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/enterprise.py @@ -0,0 +1,234 @@ +"""Enterprise Gateway adapter for the working-environment protocol. + +The existing Gateway can reconnect to and delete an already allocated sandbox, +but it cannot materialize a Home Snapshot or create a protocol-compliant +Binding. One physical sandbox owns both the materialized Home and Workspace, so +their cleanup is coupled. Runtime access remains operation-local and is routed +through the Gateway's shellctl proxy. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import logging +import shlex +from typing import cast +from urllib.parse import quote + +import httpx2 as httpx + +from dify_agent.adapters.shell.protocols import ShellCommandProtocol, ShellProviderError +from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol, ShellctlCommands +from dify_agent.runtime_backend.errors import ( + BindingAcquireError, + BindingDestroyError, + BindingLostError, + WorkspacePreservationUnsupportedError, +) +from dify_agent.runtime_backend.protocols import ( + ExecutionBindingAllocation, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + FileSystem, + HomeSnapshotCreateSpec, + InitializeHomeSnapshotSpec, + RuntimeLayout, + RuntimeLease, +) +from dify_agent.runtime_backend.shellctl import ( + ShellctlRuntimeLease, + create_owned_shellctl_lease, + run_shellctl_control_command, +) + +logger = logging.getLogger(__name__) + + +def _not_implemented() -> NotImplementedError: + return NotImplementedError("Enterprise Gateway does not implement the Execution Binding protocol") + + +@dataclass(slots=True) +class EnterpriseHomeSnapshotBackend: + """Reject Home Snapshot operations until the Gateway exposes immutable snapshots.""" + + gateway_endpoint: str + auth_token: str + gateway_timeout: float = 30.0 + + async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str: + del spec + raise _not_implemented() + + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + del spec, source + raise _not_implemented() + + async def delete(self, snapshot_ref: str) -> None: + del snapshot_ref + raise _not_implemented() + + +@dataclass(slots=True) +class EnterpriseExecutionBindingBackend: + """Access and destroy legacy Gateway sandboxes as coupled physical Bindings.""" + + gateway_endpoint: str + auth_token: str + gateway_timeout: float = 30.0 + proxy_timeout: float = 60.0 + layout: RuntimeLayout = field( + default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace") + ) + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + del spec + raise _not_implemented() + + async def acquire(self, binding_ref: str) -> RuntimeLease: + """Reconnect to one existing Gateway sandbox without creating a replacement.""" + data_plane: ShellctlRuntimeLease | None = None + try: + data_plane = await self._create_data_plane(binding_ref) + validation_commands = ShellctlCommands(client=data_plane.client) + result = await run_shellctl_control_command( + validation_commands, + "\n".join( + [ + "set -eu", + f"test -d {shlex.quote(self.layout.home_dir)}", + f"test -d {shlex.quote(self.layout.workspace_dir)}", + ] + ), + timeout=5.0, + ) + if result.exit_code != 0: + raise BindingLostError(f"Enterprise Binding {binding_ref!r} no longer contains its Home or Workspace") + return EnterpriseRuntimeLease(data_plane=data_plane) + except ShellProviderError as exc: + await _close_best_effort(data_plane, binding_ref=binding_ref) + if _is_missing_sandbox(exc): + raise BindingLostError(f"Enterprise Binding {binding_ref!r} no longer exists") from exc + raise BindingAcquireError(str(exc)) from exc + except BindingLostError: + await _close_best_effort(data_plane, binding_ref=binding_ref) + raise + except BaseException as exc: + await _close_best_effort(data_plane, binding_ref=binding_ref) + if isinstance(exc, Exception): + raise BindingAcquireError(str(exc)) from exc + raise + + async def release(self, lease: RuntimeLease) -> None: + """Close operation-local shellctl resources without deleting the sandbox.""" + if not isinstance(lease, EnterpriseRuntimeLease): + raise TypeError("EnterpriseExecutionBindingBackend can only release its own RuntimeLease") + try: + await lease.data_plane.close() + except Exception as exc: + raise BindingAcquireError(str(exc)) from exc + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + """Delete the coupled sandbox only when its Workspace is also retired.""" + if not spec.destroy_workspace: + raise WorkspacePreservationUnsupportedError( + "current Enterprise backend cannot destroy a Binding while preserving its Workspace" + ) + if spec.workspace_ref != spec.binding_ref: + raise BindingDestroyError("Enterprise Workspace ref must equal its Binding ref") + + headers = {"X-Inner-Api-Key": self.auth_token} if self.auth_token else {} + encoded_binding_ref = quote(spec.binding_ref, safe="") + try: + async with httpx.AsyncClient( + base_url=self.gateway_endpoint.rstrip("/"), + headers=headers, + timeout=httpx.Timeout(self.gateway_timeout), + ) as client: + response = await client.delete(f"/v1/sandboxes/{encoded_binding_ref}") + if response.status_code == 404: + return + _ = response.raise_for_status() + except (httpx.TimeoutException, httpx.RequestError, httpx.HTTPStatusError) as exc: + raise BindingDestroyError(str(exc)) from exc + + async def _create_data_plane(self, binding_ref: str) -> ShellctlRuntimeLease: + proxy_base_url = f"{self.gateway_endpoint.rstrip('/')}/proxy/" + headers = {"X-Sandbox-Id": binding_ref} + if self.auth_token: + headers["X-Inner-Api-Key"] = self.auth_token + http_client = httpx.AsyncClient( + base_url=proxy_base_url, + headers=headers, + follow_redirects=True, + timeout=httpx.Timeout(self.proxy_timeout), + transport=httpx.AsyncHTTPTransport(retries=3), + ) + + def client_factory() -> ShellctlClientProtocol: + from shellctl.client import ShellctlClient + + return cast( + ShellctlClientProtocol, + cast(object, ShellctlClient(proxy_base_url, token=self.auth_token, client=http_client)), + ) + + return await create_owned_shellctl_lease( + handle=binding_ref, + layout=self.layout, + entrypoint=proxy_base_url, + token=self.auth_token, + client_factory=client_factory, + owned_transport=http_client, + ) + + +@dataclass(slots=True) +class EnterpriseRuntimeLease: + """Invocation-local Enterprise shellctl connection and canonical layout.""" + + data_plane: ShellctlRuntimeLease + + @property + def handle(self) -> str: + return self.data_plane.handle + + @property + def layout(self) -> RuntimeLayout: + return self.data_plane.layout + + @property + def commands(self) -> ShellCommandProtocol: + return self.data_plane.commands + + @property + def files(self) -> FileSystem: + return self.data_plane.files + + +def _is_missing_sandbox(exc: ShellProviderError) -> bool: + return exc.status_code == 404 or (exc.code or "").casefold() in { + "not_found", + "sandbox_expired", + "sandbox_not_found", + } + + +async def _close_best_effort(data_plane: ShellctlRuntimeLease | None, *, binding_ref: str) -> None: + if data_plane is None: + return + try: + await data_plane.close() + except BaseException: + logger.warning( + "failed to close Enterprise RuntimeLease after acquisition failed", + exc_info=True, + extra={"binding_ref": binding_ref}, + ) + + +__all__ = [ + "EnterpriseExecutionBindingBackend", + "EnterpriseHomeSnapshotBackend", + "EnterpriseRuntimeLease", +] diff --git a/dify-agent/src/dify_agent/runtime_backend/errors.py b/dify-agent/src/dify_agent/runtime_backend/errors.py new file mode 100644 index 00000000000..b7aa48b6fcd --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/errors.py @@ -0,0 +1,73 @@ +"""Stable failures exposed above provider-specific runtime backends.""" + + +class RuntimeBackendError(RuntimeError): + """Base infrastructure failure for the selected runtime backend.""" + + +class HomeSnapshotCreateError(RuntimeBackendError): + pass + + +class HomeSnapshotNotFoundError(RuntimeBackendError): + pass + + +class BindingCreateError(RuntimeBackendError): + pass + + +class BindingAcquireError(RuntimeBackendError): + pass + + +class BindingLostError(RuntimeBackendError): + pass + + +class BindingDestroyError(RuntimeBackendError): + pass + + +class SharedWorkspaceUnsupportedError(RuntimeBackendError): + pass + + +class WorkspacePreservationUnsupportedError(RuntimeBackendError): + pass + + +class WorkspaceUnavailableError(RuntimeBackendError): + pass + + +class WorkspacePathError(RuntimeBackendError): + pass + + +class WorkspaceFileTooLargeError(RuntimeBackendError): + path: str + size: int + max_bytes: int + + def __init__(self, *, path: str, size: int, max_bytes: int) -> None: + self.path = path + self.size = size + self.max_bytes = max_bytes + super().__init__(f"Workspace file {path!r} exceeds the {max_bytes}-byte ToolFile upload limit") + + +__all__ = [ + "BindingAcquireError", + "BindingCreateError", + "BindingDestroyError", + "BindingLostError", + "HomeSnapshotCreateError", + "HomeSnapshotNotFoundError", + "RuntimeBackendError", + "SharedWorkspaceUnsupportedError", + "WorkspaceFileTooLargeError", + "WorkspacePathError", + "WorkspacePreservationUnsupportedError", + "WorkspaceUnavailableError", +] diff --git a/dify-agent/src/dify_agent/runtime_backend/leases.py b/dify-agent/src/dify_agent/runtime_backend/leases.py new file mode 100644 index 00000000000..82f9efed92e --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/leases.py @@ -0,0 +1,35 @@ +"""Operation-scoped RuntimeLease acquisition and release.""" + +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +import logging + +from dify_agent.runtime_backend.protocols import ExecutionBindingBackend, RuntimeLease + +logger = logging.getLogger(__name__) + + +@asynccontextmanager +async def open_runtime_lease( + backend: ExecutionBindingBackend, + binding_ref: str, +) -> AsyncGenerator[RuntimeLease, None]: + """Acquire one Binding and deterministically release its lease.""" + + lease = await backend.acquire(binding_ref) + primary_error: BaseException | None = None + try: + yield lease + except BaseException as exc: + primary_error = exc + raise + finally: + try: + await backend.release(lease) + except BaseException: + if primary_error is None: + raise + logger.warning("failed to release RuntimeLease after operation failed", exc_info=True) + + +__all__ = ["open_runtime_lease"] diff --git a/dify-agent/src/dify_agent/runtime_backend/local.py b/dify-agent/src/dify_agent/runtime_backend/local.py new file mode 100644 index 00000000000..45f7d6b8345 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/local.py @@ -0,0 +1,331 @@ +"""Local working-environment backend over one shared shellctl daemon. + +Materialized Homes and Workspaces live under separate roots. A Binding ref +encodes both logical path segments so this stateless adapter can reacquire the +same pair without maintaining a resource catalog. +""" + +from __future__ import annotations + +import logging +import posixpath +import re +import shlex +from dataclasses import dataclass + +from dify_agent.adapters.shell.protocols import ShellCommandProtocol +from dify_agent.adapters.shell.shellctl import ShellctlClientFactory +from dify_agent.runtime_backend.errors import ( + BindingAcquireError, + BindingCreateError, + BindingDestroyError, + BindingLostError, + HomeSnapshotCreateError, +) +from dify_agent.runtime_backend.protocols import ( + ExecutionBindingAllocation, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotCreateSpec, + InitializeHomeSnapshotSpec, + RuntimeLayout, + RuntimeLease, +) +from dify_agent.runtime_backend.shellctl import ( + ShellctlRuntimeLease, + create_shellctl_lease, + run_shellctl_control_command, +) + +_SAFE_REF_PART = re.compile(r"^[A-Za-z0-9._-]+$") +_BINDING_REF_SEPARATOR = ":" +logger = logging.getLogger(__name__) + + +@dataclass(slots=True) +class LocalHomeSnapshotBackend: + endpoint: str + auth_token: str + snapshot_root: str = "/home/dify/.dify-agent-home-snapshots" + client_factory: ShellctlClientFactory | None = None + + async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str: + snapshot_ref = _local_snapshot_ref(spec.home_snapshot_id) + lease = self._control_lease(snapshot_ref) + target = self._snapshot_dir(snapshot_ref) + try: + result = await run_shellctl_control_command( + lease.commands, + f"set -eu\nmkdir -p {shlex.quote(target)}\nchmod 700 {shlex.quote(target)}", + ) + if result.exit_code != 0: + raise HomeSnapshotCreateError(result.output) + return snapshot_ref + except BaseException as exc: + await _remove_partial(lease.commands, target=target, resource_ref=snapshot_ref) + if isinstance(exc, HomeSnapshotCreateError): + raise + if isinstance(exc, Exception): + raise HomeSnapshotCreateError(str(exc)) from exc + raise + finally: + await _close_best_effort(lease, resource_ref=snapshot_ref) + + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + snapshot_ref = _local_snapshot_ref(spec.home_snapshot_id) + target = self._snapshot_dir(snapshot_ref) + lease = self._control_lease(snapshot_ref, paths=(source.layout.home_dir, target)) + script = "\n".join( + [ + "set -eu", + f"test -d {shlex.quote(source.layout.home_dir)}", + f"mkdir -p {shlex.quote(target)}", + f"cp -a {shlex.quote(source.layout.home_dir)}/. {shlex.quote(target)}/", + f"chmod 700 {shlex.quote(target)}", + ] + ) + try: + result = await run_shellctl_control_command(lease.commands, script) + if result.exit_code != 0: + raise HomeSnapshotCreateError(result.output) + return snapshot_ref + except BaseException as exc: + await _remove_partial(lease.commands, target=target, resource_ref=snapshot_ref) + if isinstance(exc, HomeSnapshotCreateError): + raise + if isinstance(exc, Exception): + raise HomeSnapshotCreateError(str(exc)) from exc + raise + finally: + await _close_best_effort(lease, resource_ref=snapshot_ref) + + async def delete(self, snapshot_ref: str) -> None: + normalized = _validated_ref_part(snapshot_ref) + lease = self._control_lease(normalized) + try: + result = await run_shellctl_control_command( + lease.commands, + f"rm -rf -- {shlex.quote(self._snapshot_dir(normalized))}", + ) + if result.exit_code != 0: + raise BindingDestroyError(result.output) + except BaseException: + await _close_best_effort(lease, resource_ref=normalized) + raise + else: + await lease.close() + + def _snapshot_dir(self, snapshot_ref: str) -> str: + return f"{self.snapshot_root.rstrip('/')}/{snapshot_ref}" + + def _control_lease(self, handle: str, *, paths: tuple[str, ...] = ()) -> ShellctlRuntimeLease: + control_root = _control_root(paths or (self.snapshot_root,)) + layout = RuntimeLayout(home_dir=control_root, workspace_dir=control_root) + return create_shellctl_lease( + handle=handle, + layout=layout, + entrypoint=self.endpoint, + token=self.auth_token, + client_factory=self.client_factory, + ) + + +@dataclass(slots=True) +class LocalExecutionBindingBackend: + endpoint: str + auth_token: str + materialized_home_root: str = "/home/dify/.dify-agent-materialized-homes" + workspace_root: str = "/home/dify/.dify-agent-workspaces" + snapshot_root: str = "/home/dify/.dify-agent-home-snapshots" + client_factory: ShellctlClientFactory | None = None + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + binding_id = _validated_ref_part(spec.binding_id) + workspace_id = _validated_ref_part(spec.workspace_id) + snapshot_ref = _validated_ref_part(spec.home_snapshot_ref) + workspace_ref = workspace_id + if spec.existing_workspace_ref is not None: + existing_workspace_ref = _validated_ref_part(spec.existing_workspace_ref) + if existing_workspace_ref != workspace_ref: + raise BindingCreateError("existing Workspace ref does not match workspace_id") + binding_ref = _local_binding_ref(binding_id=binding_id, workspace_id=workspace_id) + lease = self._control_lease(binding_ref) + home_dir = self._home_dir(binding_id) + workspace_dir = self._workspace_dir(workspace_id) + snapshot_dir = f"{self.snapshot_root.rstrip('/')}/{snapshot_ref}" + creates_workspace = spec.existing_workspace_ref is None + workspace_setup = ( + f"mkdir -p {shlex.quote(workspace_dir)}" if creates_workspace else f"test -d {shlex.quote(workspace_dir)}" + ) + script = "\n".join( + [ + "set -eu", + f"test -d {shlex.quote(snapshot_dir)}", + workspace_setup, + f"mkdir -p {shlex.quote(home_dir)}", + f"cp -a {shlex.quote(snapshot_dir)}/. {shlex.quote(home_dir)}/", + f"chmod 700 {shlex.quote(home_dir)} {shlex.quote(workspace_dir)}", + ] + ) + try: + result = await run_shellctl_control_command(lease.commands, script) + if result.exit_code != 0: + raise BindingCreateError(result.output) + return ExecutionBindingAllocation(binding_ref=binding_ref, workspace_ref=workspace_ref) + except BaseException as exc: + targets = [home_dir] + if creates_workspace: + targets.append(workspace_dir) + await _remove_partial( + lease.commands, + target=" ".join(shlex.quote(target) for target in targets), + resource_ref=binding_ref, + target_is_shell_words=True, + ) + if isinstance(exc, BindingCreateError): + raise + if isinstance(exc, Exception): + raise BindingCreateError(str(exc)) from exc + raise + finally: + await _close_best_effort(lease, resource_ref=binding_ref) + + async def acquire(self, binding_ref: str) -> RuntimeLease: + lease = self._lease(binding_ref) + try: + result = await run_shellctl_control_command( + lease.commands, + "\n".join( + [ + "set -eu", + f"test -d {shlex.quote(lease.layout.home_dir)}", + f"test -d {shlex.quote(lease.layout.workspace_dir)}", + ] + ), + timeout=10.0, + ) + if result.exit_code != 0: + raise BindingLostError(f"Local Binding {binding_ref!r} no longer exists") + return lease + except BaseException as exc: + await _close_best_effort(lease, resource_ref=binding_ref) + if isinstance(exc, BindingLostError): + raise + if isinstance(exc, Exception): + raise BindingAcquireError(str(exc)) from exc + raise + + async def release(self, lease: RuntimeLease) -> None: + if not isinstance(lease, ShellctlRuntimeLease): + raise TypeError("LocalExecutionBindingBackend can only release its own RuntimeLease") + try: + await lease.close() + except Exception as exc: + raise BindingAcquireError(str(exc)) from exc + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + binding_id, workspace_id = _parse_local_binding_ref(spec.binding_ref) + if spec.destroy_workspace: + workspace_ref = _validated_ref_part(spec.workspace_ref or "") + if workspace_ref != workspace_id: + raise BindingDestroyError("Workspace ref does not match Binding ref") + lease = self._control_lease(spec.binding_ref) + targets = [self._home_dir(binding_id)] + if spec.destroy_workspace: + targets.append(self._workspace_dir(workspace_id)) + try: + result = await run_shellctl_control_command( + lease.commands, + "rm -rf -- " + " ".join(shlex.quote(target) for target in targets), + ) + if result.exit_code != 0: + raise BindingDestroyError(result.output) + except BaseException: + await _close_best_effort(lease, resource_ref=spec.binding_ref) + raise + else: + await lease.close() + + def _lease(self, binding_ref: str) -> ShellctlRuntimeLease: + binding_id, workspace_id = _parse_local_binding_ref(binding_ref) + return create_shellctl_lease( + handle=binding_ref, + layout=RuntimeLayout( + home_dir=self._home_dir(binding_id), + workspace_dir=self._workspace_dir(workspace_id), + ), + entrypoint=self.endpoint, + token=self.auth_token, + client_factory=self.client_factory, + ) + + def _control_lease(self, handle: str) -> ShellctlRuntimeLease: + control_root = _control_root((self.materialized_home_root, self.workspace_root, self.snapshot_root)) + return create_shellctl_lease( + handle=handle, + layout=RuntimeLayout(home_dir=control_root, workspace_dir=control_root), + entrypoint=self.endpoint, + token=self.auth_token, + client_factory=self.client_factory, + ) + + def _home_dir(self, binding_id: str) -> str: + return f"{self.materialized_home_root.rstrip('/')}/{binding_id}" + + def _workspace_dir(self, workspace_id: str) -> str: + return f"{self.workspace_root.rstrip('/')}/{workspace_id}" + + +def _local_snapshot_ref(home_snapshot_id: str) -> str: + return f"home-{_validated_ref_part(home_snapshot_id)}" + + +def _local_binding_ref(*, binding_id: str, workspace_id: str) -> str: + return f"{_validated_ref_part(binding_id)}{_BINDING_REF_SEPARATOR}{_validated_ref_part(workspace_id)}" + + +def _parse_local_binding_ref(binding_ref: str) -> tuple[str, str]: + parts = binding_ref.split(_BINDING_REF_SEPARATOR) + if len(parts) != 2: + raise ValueError("Local Binding ref is invalid") + return _validated_ref_part(parts[0]), _validated_ref_part(parts[1]) + + +def _validated_ref_part(value: str) -> str: + if value in {"", ".", ".."} or _SAFE_REF_PART.fullmatch(value) is None: + raise ValueError("runtime backend ref must be a safe path segment") + return value + + +def _control_root(paths: tuple[str, ...]) -> str: + normalized = tuple(posixpath.normpath(path) for path in paths) + common = posixpath.commonpath(normalized) + if len(normalized) == 1 or common in normalized: + common = posixpath.dirname(common) + return common or "/" + + +async def _remove_partial( + commands: ShellCommandProtocol, + *, + target: str, + resource_ref: str, + target_is_shell_words: bool = False, +) -> None: + try: + target_words = target if target_is_shell_words else shlex.quote(target) + result = await run_shellctl_control_command(commands, f"rm -rf -- {target_words}") + if result.exit_code != 0: + logger.warning("failed to remove partial local resource", extra={"resource_ref": resource_ref}) + except BaseException: + logger.warning("failed to remove partial local resource", exc_info=True, extra={"resource_ref": resource_ref}) + + +async def _close_best_effort(lease: ShellctlRuntimeLease, *, resource_ref: str) -> None: + try: + await lease.close() + except BaseException: + logger.warning("failed to close local RuntimeLease", exc_info=True, extra={"resource_ref": resource_ref}) + + +__all__ = ["LocalExecutionBindingBackend", "LocalHomeSnapshotBackend"] diff --git a/dify-agent/src/dify_agent/runtime_backend/profile.py b/dify-agent/src/dify_agent/runtime_backend/profile.py new file mode 100644 index 00000000000..07441e28fb7 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/profile.py @@ -0,0 +1,137 @@ +"""Deployment-selected coherent runtime backend profile construction.""" + +from __future__ import annotations + +from typing import ClassVar, Literal, Self +from urllib.parse import urlparse + +from pydantic import AliasChoices, Field, model_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + +from dify_agent.runtime_backend.e2b import ( + E2B_MAX_ACTIVE_TIMEOUT_SECONDS, + E2BHomeSnapshotBackend, + E2BSDKControlPlane, + E2BExecutionBindingBackend, +) +from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend, EnterpriseHomeSnapshotBackend +from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend +from dify_agent.runtime_backend.protocols import RuntimeBackendProfile + +DEFAULT_E2B_TEMPLATE = "difys-default-team/dify-agent-local-sandbox" + + +class RuntimeBackendSettings(BaseSettings): + """Server-private credentials and endpoints for one coherent backend profile.""" + + runtime_backend: Literal["local", "enterprise", "e2b"] = "local" + + local_sandbox_endpoint: str | None = Field( + default=None, + validation_alias=AliasChoices( + "DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT", + "DIFY_AGENT_SHELLCTL_ENTRYPOINT", + ), + ) + local_sandbox_auth_token: str | None = Field( + default=None, + validation_alias=AliasChoices( + "DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN", + "DIFY_AGENT_SHELLCTL_AUTH_TOKEN", + ), + ) + + enterprise_sandbox_gateway_endpoint: str | None = None + enterprise_sandbox_gateway_auth_token: str | None = None + enterprise_sandbox_gateway_timeout: float = Field(default=30.0, gt=0) + enterprise_sandbox_proxy_timeout: float = Field(default=60.0, gt=0) + + e2b_api_key: str | None = None + e2b_template: str = DEFAULT_E2B_TEMPLATE + e2b_active_timeout_seconds: int = Field( + default=E2B_MAX_ACTIVE_TIMEOUT_SECONDS, + ge=1, + le=E2B_MAX_ACTIVE_TIMEOUT_SECONDS, + ) + e2b_shellctl_auth_token: str = "" + e2b_shellctl_port: int = Field(default=5004, ge=1, le=65535) + + model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict( + env_prefix="DIFY_AGENT_", + env_file=(".env", "dify-agent/.env"), + extra="ignore", + populate_by_name=True, + ) + + @model_validator(mode="after") + def validate_selected_backend(self) -> Self: + match self.runtime_backend: + case "local": + if not self.local_sandbox_endpoint or not self.local_sandbox_endpoint.strip(): + raise ValueError("local_sandbox_endpoint is required for the local runtime backend") + _validate_http_url(self.local_sandbox_endpoint, field_name="local_sandbox_endpoint") + case "enterprise": + endpoint = self.enterprise_sandbox_gateway_endpoint + if not endpoint or not endpoint.strip(): + raise ValueError( + "enterprise_sandbox_gateway_endpoint is required for the enterprise runtime backend" + ) + _validate_http_url(endpoint, field_name="enterprise_sandbox_gateway_endpoint") + case "e2b": + if not self.e2b_api_key or not self.e2b_api_key.strip(): + raise ValueError("e2b_api_key is required for the e2b runtime backend") + if not self.e2b_template.strip(): + raise ValueError("e2b_template must not be blank") + return self + + +def create_runtime_backend_profile(settings: RuntimeBackendSettings) -> RuntimeBackendProfile: + """Construct one driver pair selected exclusively by server deployment settings.""" + match settings.runtime_backend: + case "local": + endpoint = settings.local_sandbox_endpoint or "" + token = settings.local_sandbox_auth_token or "" + return RuntimeBackendProfile( + home_snapshots=LocalHomeSnapshotBackend(endpoint=endpoint, auth_token=token), + execution_bindings=LocalExecutionBindingBackend(endpoint=endpoint, auth_token=token), + ) + case "enterprise": + endpoint = settings.enterprise_sandbox_gateway_endpoint or "" + token = settings.enterprise_sandbox_gateway_auth_token or "" + return RuntimeBackendProfile( + home_snapshots=EnterpriseHomeSnapshotBackend( + gateway_endpoint=endpoint, + auth_token=token, + gateway_timeout=settings.enterprise_sandbox_gateway_timeout, + ), + execution_bindings=EnterpriseExecutionBindingBackend( + gateway_endpoint=endpoint, + auth_token=token, + gateway_timeout=settings.enterprise_sandbox_gateway_timeout, + proxy_timeout=settings.enterprise_sandbox_proxy_timeout, + ), + ) + case "e2b": + control_plane = E2BSDKControlPlane(api_key=settings.e2b_api_key or "") + return RuntimeBackendProfile( + home_snapshots=E2BHomeSnapshotBackend( + control_plane=control_plane, + template=settings.e2b_template, + active_timeout_seconds=settings.e2b_active_timeout_seconds, + ), + execution_bindings=E2BExecutionBindingBackend( + control_plane=control_plane, + active_timeout_seconds=settings.e2b_active_timeout_seconds, + shellctl_auth_token=settings.e2b_shellctl_auth_token, + shellctl_port=settings.e2b_shellctl_port, + ), + ) + + +def _validate_http_url(value: str, *, field_name: str) -> None: + parsed = urlparse(value.strip()) + if parsed.scheme not in {"http", "https"} or not parsed.netloc: + raise ValueError(f"{field_name} must be a valid http(s) URL") + + +__all__ = ["DEFAULT_E2B_TEMPLATE", "RuntimeBackendSettings", "create_runtime_backend_profile"] diff --git a/dify-agent/src/dify_agent/runtime_backend/protocols.py b/dify-agent/src/dify_agent/runtime_backend/protocols.py new file mode 100644 index 00000000000..343bea899a7 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/protocols.py @@ -0,0 +1,228 @@ +"""Backend-neutral contracts for persistent working environments. + +Dify API owns logical Home Snapshot, Workspace, and Agent Workspace Binding +records. Backends own their physical representations. ``RuntimeLease`` is the +only invocation-local object and must never be serialized or persisted. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from dify_agent.adapters.shell.protocols import ShellCommandProtocol + + +@dataclass(frozen=True, slots=True) +class InitializeHomeSnapshotSpec: + tenant_id: str + agent_id: str + home_snapshot_id: str + + +@dataclass(frozen=True, slots=True) +class HomeSnapshotCreateSpec: + tenant_id: str + agent_id: str + home_snapshot_id: str + + +@dataclass(frozen=True, slots=True) +class RuntimeLayout: + """Canonical Home and Workspace roots exposed for one operation.""" + + home_dir: str + workspace_dir: str + + +@dataclass(frozen=True, slots=True) +class WorkspaceFileEntry: + name: str + type: str + size: int | None + mtime: int | None + + +@dataclass(frozen=True, slots=True) +class WorkspaceListResult: + path: str + entries: tuple[WorkspaceFileEntry, ...] + truncated: bool + + +@dataclass(frozen=True, slots=True) +class WorkspaceReadResult: + path: str + size: int + truncated: bool + binary: bool + text: str | None + + +@dataclass(frozen=True, slots=True) +class WorkspaceFileContent: + path: str + size: int + content: bytes + + +class FileSystem(Protocol): + """File operations interpreted in the current RuntimeLease namespace.""" + + async def list_directory(self, *, path: str, limit: int) -> WorkspaceListResult: ... + + async def read_file(self, *, path: str, max_bytes: int) -> WorkspaceReadResult: ... + + async def read_bytes(self, *, path: str, max_bytes: int) -> WorkspaceFileContent: ... + + async def upload(self, *, content: bytes, remote_path: str, cwd: str | None = None) -> None: ... + + async def download(self, *, remote_path: str, cwd: str | None = None) -> bytes: ... + + +class RuntimeLease(Protocol): + """Invocation-local data-plane access to one persistent Binding.""" + + @property + def layout(self) -> RuntimeLayout: ... + + @property + def commands(self) -> ShellCommandProtocol: ... + + @property + def files(self) -> FileSystem: ... + + +@dataclass(frozen=True, slots=True) +class ExecutionBindingCreateSpec: + tenant_id: str + agent_id: str + binding_id: str + workspace_id: str + existing_workspace_ref: str | None + home_snapshot_ref: str + + +@dataclass(frozen=True, slots=True) +class ExecutionBindingAllocation: + binding_ref: str + workspace_ref: str + + +@dataclass(frozen=True, slots=True) +class ExecutionBindingDestroySpec: + binding_ref: str + destroy_workspace: bool + workspace_ref: str | None = None + + def __post_init__(self) -> None: + if self.destroy_workspace and not self.workspace_ref: + raise ValueError("workspace_ref is required when destroy_workspace is true") + + +class ExecutionBindingBackend(Protocol): + """Manage physical Binding, Materialized Home, and Workspace resources.""" + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + """Materialize a mutable Home and make the requested Workspace ready. + + Implementations must initialize the Home from ``home_snapshot_ref`` and + return stable opaque refs only after both resources are usable. With an + ``existing_workspace_ref``, they must attach that Workspace without + clearing or replacing its contents; unsupported sharing must fail before + mutating it. On failure, implementations should clean up newly allocated + partial resources and must not damage a pre-existing Workspace. + """ + ... + + async def acquire(self, binding_ref: str) -> RuntimeLease: + """Open fresh operation-scoped access to an existing physical Binding. + + Implementations must reactivate or reconnect the resource as needed, + verify that its materialized Home and Workspace still exist, and must not + silently create a replacement. Confirmed resource loss must raise + ``BindingLostError``; other acquisition failures must raise + ``BindingAcquireError``. + """ + ... + + async def release(self, lease: RuntimeLease) -> None: + """End operation-local access without destroying persistent resources. + + Implementations must close clients and owned transports and may suspend + the physical runtime. They must not delete or retire the Binding, + materialized Home, or Workspace, and must reject leases created by a + different backend implementation. + """ + ... + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + """Destroy a Binding's materialized Home and optionally its Workspace. + + Implementations must honor ``destroy_workspace`` and validate + ``workspace_ref`` before destructive work. A backend unable to preserve + the Workspace must raise ``WorkspacePreservationUnsupportedError`` before + deleting anything. Cleanup must be safe to retry and must never delete + the immutable Home Snapshot from which the Binding was materialized. + """ + ... + + +class HomeSnapshotBackend(Protocol): + """Manage immutable backend-native Home resources.""" + + async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str: + """Create the deployment-defined baseline Home Snapshot. + + Implementations may use any backend-native bootstrap mechanism, but must + return a stable opaque ref only after an immutable snapshot is ready for + future Binding creation. Temporary bootstrap resources must not become + part of the logical snapshot lifecycle. + """ + ... + + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + """Capture the source lease's current Home as a new immutable snapshot. + + Implementations must not mutate or release ``source``. Workspace content + is not logically part of a Home Snapshot; if a provider can only snapshot + a coupled runtime, later materialization must prevent captured Workspace + state from becoming the new Binding's Workspace. + """ + ... + + async def delete(self, snapshot_ref: str) -> None: + """Delete one physical snapshot without affecting derived resources. + + Implementations must make deletion safe to retry, treat an already absent + snapshot as a successful outcome, and must not delete Bindings, + materialized Homes, or Workspaces created from the snapshot. + """ + ... + + +@dataclass(frozen=True, slots=True) +class RuntimeBackendProfile: + """Coherent Home and Binding backends selected once per deployment.""" + + home_snapshots: HomeSnapshotBackend + execution_bindings: ExecutionBindingBackend + + +__all__ = [ + "ExecutionBindingAllocation", + "ExecutionBindingBackend", + "ExecutionBindingCreateSpec", + "ExecutionBindingDestroySpec", + "FileSystem", + "HomeSnapshotBackend", + "HomeSnapshotCreateSpec", + "InitializeHomeSnapshotSpec", + "RuntimeBackendProfile", + "RuntimeLayout", + "RuntimeLease", + "WorkspaceFileContent", + "WorkspaceFileEntry", + "WorkspaceListResult", + "WorkspaceReadResult", +] diff --git a/dify-agent/src/dify_agent/runtime_backend/shellctl.py b/dify-agent/src/dify_agent/runtime_backend/shellctl.py new file mode 100644 index 00000000000..257157dfc82 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/shellctl.py @@ -0,0 +1,156 @@ +"""Shared shellctl RuntimeLease used by every runtime backend.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import logging +from typing import Protocol + +from dify_agent.adapters.shell.protocols import CompleteShellCommandResult, ShellCommandProtocol +from dify_agent.adapters.shell.shellctl import ( + ShellctlClientFactory, + ShellctlClientProtocol, + ShellctlCommands, + ShellctlFileTransfer, + create_default_shellctl_client_factory, +) +from dify_agent.runtime_backend.protocols import FileSystem, RuntimeLayout + +_CONTROL_COMMAND_OUTPUT_LIMIT = 256 * 1024 +logger = logging.getLogger(__name__) + + +class AsyncCloseable(Protocol): + async def aclose(self) -> None: ... + + +@dataclass(slots=True) +class ShellctlRuntimeLease: + """One invocation-local shellctl connection and its canonical layout.""" + + handle: str + layout: RuntimeLayout + client: ShellctlClientProtocol + commands: ShellCommandProtocol + files: FileSystem + owned_transport: AsyncCloseable | None = None + _closed: bool = field(default=False, init=False) + + async def close(self) -> None: + if self._closed: + return + self._closed = True + client_error: BaseException | None = None + try: + await self.client.close() + except BaseException as exc: + client_error = exc + try: + if self.owned_transport is not None: + await self.owned_transport.aclose() + except BaseException as exc: + if client_error is None: + raise + logger.warning("Failed to close owned shellctl transport after client close failed: %s", exc) + if client_error is not None: + raise client_error + + +def create_shellctl_lease( + *, + handle: str, + layout: RuntimeLayout, + entrypoint: str, + token: str, + client_factory: ShellctlClientFactory | None = None, + owned_transport: AsyncCloseable | None = None, +) -> ShellctlRuntimeLease: + """Create adapters around one new shellctl client without owning control-plane lifecycle.""" + factory = client_factory or create_default_shellctl_client_factory(entrypoint=entrypoint, token=token) + client = factory() + return ShellctlRuntimeLease( + handle=handle, + layout=layout, + client=client, + commands=ShellctlCommands( + client=client, + home_dir=layout.home_dir, + workspace_dir=layout.workspace_dir, + ), + files=ShellctlFileTransfer( + client=client, + cwd=layout.workspace_dir, + home_dir=layout.home_dir, + ), + owned_transport=owned_transport, + ) + + +async def create_owned_shellctl_lease( + *, + handle: str, + layout: RuntimeLayout, + entrypoint: str, + token: str, + client_factory: ShellctlClientFactory, + owned_transport: AsyncCloseable, +) -> ShellctlRuntimeLease: + """Create a lease that owns an injected transport, closing it if construction fails.""" + try: + return create_shellctl_lease( + handle=handle, + layout=layout, + entrypoint=entrypoint, + token=token, + client_factory=client_factory, + owned_transport=owned_transport, + ) + except BaseException: + try: + await owned_transport.aclose() + except BaseException as cleanup_exc: + logger.warning("Failed to close owned shellctl transport after lease construction failed: %s", cleanup_exc) + raise + + +async def run_shellctl_control_command( + commands: ShellCommandProtocol, + script: str, + *, + timeout: float = 30.0, +) -> CompleteShellCommandResult: + """Run one bounded driver control command and always delete its transient job.""" + result = await commands.run(script, cwd=None, env=None, timeout=timeout) + job_id = result.job_id + output_parts = [result.output] + try: + while result.truncated or not result.done: + result = await commands.wait(result.job_id, offset=result.offset, timeout=timeout) + output_parts.append(result.output) + if sum(len(part.encode("utf-8")) for part in output_parts) > _CONTROL_COMMAND_OUTPUT_LIMIT: + raise RuntimeError("shellctl control command exceeded its output limit") + return CompleteShellCommandResult( + job_id=result.job_id, + status=result.status, + done=result.done, + exit_code=result.exit_code, + output="".join(output_parts), + output_complete=True, + incomplete_reason=None, + offset=result.offset, + output_path=result.output_path, + ) + finally: + try: + await commands.delete(job_id, force=True) + except Exception as exc: + logger.warning("Failed to delete transient shellctl control job %s: %s", job_id, exc) + + +__all__ = [ + "AsyncCloseable", + "ShellctlRuntimeLease", + "create_owned_shellctl_lease", + "create_shellctl_lease", + "run_shellctl_control_command", +] diff --git a/dify-agent/src/dify_agent/server/app.py b/dify-agent/src/dify_agent/server/app.py index 1fadd208d57..e19af9a75ad 100644 --- a/dify-agent/src/dify_agent/server/app.py +++ b/dify-agent/src/dify_agent/server/app.py @@ -7,8 +7,8 @@ rather than request handlers, so client disconnects do not cancel the agent runtime. Redis persists run records and per-run event streams with configured retention only; it is not used as a job queue. Agenton layers and providers stay state-only: they borrow the lifespan-owned clients through the runner and -receive shell-layer server settings through provider construction rather than -reading environment variables themselves. The standard server always mounts the +receive runtime-backend and Shell settings through provider construction rather +than reading environment variables themselves. The standard server always mounts the HTTP Agent Stub router and additionally starts the optional grpclib Agent Stub server when ``DIFY_AGENT_STUB_API_BASE_URL`` uses ``grpc://``. Process-level Logfire instrumentation is configured at app construction time and only exports @@ -32,8 +32,12 @@ from dify_agent.runtime.run_scheduler import RunScheduler from dify_agent.server.auth import create_bearer_token_dependency from dify_agent.server.observability import configure_server_observability from dify_agent.server.routes.runs import create_runs_router -from dify_agent.server.routes.sandbox_files import create_sandbox_files_router -from dify_agent.server.sandbox_files import SandboxFileService +from dify_agent.server.routes.execution_bindings import create_execution_bindings_router +from dify_agent.server.routes.home_snapshots import create_home_snapshots_router +from dify_agent.server.routes.workspace_files import create_workspace_files_router +from dify_agent.server.execution_bindings import ExecutionBindingService +from dify_agent.server.workspace_files import AgentStubWorkspaceFileUploader, WorkspaceFileService +from dify_agent.server.home_snapshots import HomeSnapshotService from dify_agent.server.settings import ServerSettings from dify_agent.storage.redis_run_store import RedisRunStore @@ -60,19 +64,43 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI: agent_stub_file_request_handler = resolved_settings.create_agent_stub_file_request_handler() agent_stub_config_request_handler = resolved_settings.create_agent_stub_config_request_handler() agent_stub_drive_request_handler = resolved_settings.create_agent_stub_drive_request_handler() - shell_provider = resolved_settings.build_shell_provider() + runtime_backend_profile = resolved_settings.build_runtime_backend_profile() layer_providers = create_default_layer_providers( plugin_daemon_url=resolved_settings.plugin_daemon_url, plugin_daemon_api_key=resolved_settings.plugin_daemon_api_key, inner_api_url=resolved_settings.inner_api_url, inner_api_key=resolved_settings.inner_api_key or "", - shell_provider=shell_provider, - shell_home_root=resolved_settings.shell_home_root, + runtime_backend_profile=runtime_backend_profile, shell_redact_patterns=resolved_settings.get_shell_redact_patterns(), agent_stub_api_base_url=resolved_settings.agent_stub_api_base_url, agent_stub_token_factory=agent_stub_token_factory, ) - sandbox_file_service = SandboxFileService(layer_providers=layer_providers) if shell_provider is not None else None + workspace_file_service = ( + WorkspaceFileService( + execution_bindings=runtime_backend_profile.execution_bindings, + upload_max_bytes=resolved_settings.sandbox_file_upload_max_bytes, + file_uploader=( + AgentStubWorkspaceFileUploader(file_request_handler=agent_stub_file_request_handler) + if agent_stub_file_request_handler is not None + else None + ), + ) + if runtime_backend_profile is not None + else None + ) + home_snapshot_service = ( + HomeSnapshotService( + home_snapshots=runtime_backend_profile.home_snapshots, + execution_bindings=runtime_backend_profile.execution_bindings, + ) + if runtime_backend_profile is not None + else None + ) + execution_binding_service = ( + ExecutionBindingService(backend=runtime_backend_profile.execution_bindings) + if runtime_backend_profile is not None + else None + ) state: dict[str, object] = {} @asynccontextmanager @@ -131,7 +159,9 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI: auth_dependency=create_bearer_token_dependency(resolved_settings.api_token), ) ) - app.include_router(create_sandbox_files_router(lambda: sandbox_file_service)) + app.include_router(create_execution_bindings_router(lambda: execution_binding_service)) + app.include_router(create_home_snapshots_router(lambda: home_snapshot_service)) + app.include_router(create_workspace_files_router(lambda: workspace_file_service)) app.include_router( create_agent_stub_router( token_codec=agent_stub_token_codec, diff --git a/dify-agent/src/dify_agent/server/execution_bindings.py b/dify-agent/src/dify_agent/server/execution_bindings.py new file mode 100644 index 00000000000..fbecb129daa --- /dev/null +++ b/dify-agent/src/dify_agent/server/execution_bindings.py @@ -0,0 +1,73 @@ +"""Stateless application facade for Execution Binding lifecycle operations.""" + +from dataclasses import dataclass + +from dify_agent.protocol import ( + CreateExecutionBindingRequest, + CreateExecutionBindingResponse, + DestroyExecutionBindingRequest, +) +from dify_agent.runtime_backend import ( + BindingCreateError, + BindingDestroyError, + ExecutionBindingBackend, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + SharedWorkspaceUnsupportedError, + WorkspacePreservationUnsupportedError, +) + + +class ExecutionBindingServiceError(RuntimeError): + code: str + message: str + status_code: int + + def __init__(self, code: str, message: str, *, status_code: int) -> None: + super().__init__(message) + self.code = code + self.message = message + self.status_code = status_code + + +@dataclass(slots=True) +class ExecutionBindingService: + backend: ExecutionBindingBackend + + async def create_binding(self, request: CreateExecutionBindingRequest) -> CreateExecutionBindingResponse: + try: + allocation = await self.backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id=request.tenant_id, + agent_id=request.agent_id, + binding_id=request.binding_id, + workspace_id=request.workspace_id, + existing_workspace_ref=request.existing_workspace_ref, + home_snapshot_ref=request.home_snapshot_ref, + ) + ) + except SharedWorkspaceUnsupportedError as exc: + raise ExecutionBindingServiceError("shared_workspace_unsupported", str(exc), status_code=409) from exc + except BindingCreateError as exc: + raise ExecutionBindingServiceError("binding_create_failed", str(exc), status_code=502) from exc + return CreateExecutionBindingResponse( + binding_ref=allocation.binding_ref, + workspace_ref=allocation.workspace_ref, + ) + + async def destroy_binding(self, request: DestroyExecutionBindingRequest) -> None: + try: + await self.backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=request.binding_ref, + destroy_workspace=request.destroy_workspace, + workspace_ref=request.workspace_ref, + ) + ) + except WorkspacePreservationUnsupportedError as exc: + raise ExecutionBindingServiceError("workspace_preservation_unsupported", str(exc), status_code=409) from exc + except BindingDestroyError as exc: + raise ExecutionBindingServiceError("binding_destroy_failed", str(exc), status_code=502) from exc + + +__all__ = ["ExecutionBindingService", "ExecutionBindingServiceError"] diff --git a/dify-agent/src/dify_agent/server/home_snapshots.py b/dify-agent/src/dify_agent/server/home_snapshots.py new file mode 100644 index 00000000000..479a915667d --- /dev/null +++ b/dify-agent/src/dify_agent/server/home_snapshots.py @@ -0,0 +1,83 @@ +"""Stateless application facade for deployment-owned Home Snapshots.""" + +from dataclasses import dataclass + +from dify_agent.protocol.home_snapshot import ( + CreateHomeSnapshotFromBindingRequest, + DeleteHomeSnapshotRequest, + HomeSnapshotResponse, + InitializeHomeSnapshotRequest, +) +from dify_agent.runtime_backend import ( + BindingAcquireError, + BindingLostError, + ExecutionBindingBackend, + HomeSnapshotBackend, + HomeSnapshotCreateError, + HomeSnapshotCreateSpec, + HomeSnapshotNotFoundError, + InitializeHomeSnapshotSpec, +) +from dify_agent.runtime_backend.leases import open_runtime_lease + + +class HomeSnapshotServiceError(RuntimeError): + code: str + message: str + status_code: int + + def __init__(self, code: str, message: str, *, status_code: int) -> None: + super().__init__(message) + self.code = code + self.message = message + self.status_code = status_code + + +@dataclass(slots=True) +class HomeSnapshotService: + home_snapshots: HomeSnapshotBackend + execution_bindings: ExecutionBindingBackend + + async def initialize(self, request: InitializeHomeSnapshotRequest) -> HomeSnapshotResponse: + try: + snapshot_ref = await self.home_snapshots.initialize( + InitializeHomeSnapshotSpec( + tenant_id=request.tenant_id, + agent_id=request.agent_id, + home_snapshot_id=request.home_snapshot_id, + ) + ) + except HomeSnapshotCreateError as exc: + raise HomeSnapshotServiceError("home_snapshot_create_failed", str(exc), status_code=502) from exc + return HomeSnapshotResponse(snapshot_ref=snapshot_ref) + + async def create_from_binding( + self, + request: CreateHomeSnapshotFromBindingRequest, + ) -> HomeSnapshotResponse: + try: + async with open_runtime_lease(self.execution_bindings, request.backend_binding_ref) as lease: + snapshot_ref = await self.home_snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec( + tenant_id=request.tenant_id, + agent_id=request.agent_id, + home_snapshot_id=request.home_snapshot_id, + ), + source=lease, + ) + except BindingLostError as exc: + raise HomeSnapshotServiceError("binding_lost", str(exc), status_code=404) from exc + except (BindingAcquireError, HomeSnapshotCreateError) as exc: + raise HomeSnapshotServiceError("home_snapshot_create_failed", str(exc), status_code=502) from exc + return HomeSnapshotResponse(snapshot_ref=snapshot_ref) + + async def delete(self, request: DeleteHomeSnapshotRequest) -> None: + try: + await self.home_snapshots.delete(request.snapshot_ref) + except HomeSnapshotNotFoundError: + return + except RuntimeError as exc: + raise HomeSnapshotServiceError("home_snapshot_delete_failed", str(exc), status_code=502) from exc + + +__all__ = ["HomeSnapshotService", "HomeSnapshotServiceError"] diff --git a/dify-agent/src/dify_agent/server/routes/execution_bindings.py b/dify-agent/src/dify_agent/server/routes/execution_bindings.py new file mode 100644 index 00000000000..607574a36e1 --- /dev/null +++ b/dify-agent/src/dify_agent/server/routes/execution_bindings.py @@ -0,0 +1,55 @@ +"""Private Execution Binding control-plane routes used by Dify API.""" + +from collections.abc import Callable +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException, Response, status + +from dify_agent.protocol import ( + CreateExecutionBindingRequest, + CreateExecutionBindingResponse, + DestroyExecutionBindingRequest, +) +from dify_agent.server.execution_bindings import ExecutionBindingService, ExecutionBindingServiceError + + +def create_execution_bindings_router(get_service: Callable[[], ExecutionBindingService | None]) -> APIRouter: + router = APIRouter(prefix="/execution-bindings", tags=["execution-bindings"]) + + def service_dep() -> ExecutionBindingService: + service = get_service() + if service is None: + raise HTTPException( + status_code=503, + detail={"code": "runtime_backend_unavailable", "message": "runtime backend is not configured"}, + ) + return service + + def raise_http(exc: ExecutionBindingServiceError) -> HTTPException: + return HTTPException(status_code=exc.status_code, detail={"code": exc.code, "message": exc.message}) + + @router.post("", response_model=CreateExecutionBindingResponse, status_code=status.HTTP_201_CREATED) + async def create_binding( + request: CreateExecutionBindingRequest, + service: Annotated[ExecutionBindingService, Depends(service_dep)], + ) -> CreateExecutionBindingResponse: + try: + return await service.create_binding(request) + except ExecutionBindingServiceError as exc: + raise raise_http(exc) from exc + + @router.post("/destroy", status_code=status.HTTP_204_NO_CONTENT) + async def destroy_binding( + request: DestroyExecutionBindingRequest, + service: Annotated[ExecutionBindingService, Depends(service_dep)], + ) -> Response: + try: + await service.destroy_binding(request) + except ExecutionBindingServiceError as exc: + raise raise_http(exc) from exc + return Response(status_code=status.HTTP_204_NO_CONTENT) + + return router + + +__all__ = ["create_execution_bindings_router"] diff --git a/dify-agent/src/dify_agent/server/routes/home_snapshots.py b/dify-agent/src/dify_agent/server/routes/home_snapshots.py new file mode 100644 index 00000000000..52a00dafb42 --- /dev/null +++ b/dify-agent/src/dify_agent/server/routes/home_snapshots.py @@ -0,0 +1,72 @@ +"""Private Home Snapshot control-plane routes used by Dify API.""" + +from collections.abc import Callable +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException, Response, status + +from dify_agent.protocol.home_snapshot import ( + CreateHomeSnapshotFromBindingRequest, + DeleteHomeSnapshotRequest, + HomeSnapshotResponse, + InitializeHomeSnapshotRequest, +) +from dify_agent.server.home_snapshots import HomeSnapshotService, HomeSnapshotServiceError + + +def create_home_snapshots_router(get_service: Callable[[], HomeSnapshotService | None]) -> APIRouter: + router = APIRouter(prefix="/home-snapshots", tags=["home-snapshots"]) + + def service_dep() -> HomeSnapshotService: + service = get_service() + if service is None: + raise HTTPException( + status_code=503, + detail={"code": "runtime_backend_unavailable", "message": "runtime backend is not configured"}, + ) + return service + + @router.post("/initialize", response_model=HomeSnapshotResponse, status_code=status.HTTP_201_CREATED) + async def initialize_snapshot( + request: InitializeHomeSnapshotRequest, + service: Annotated[HomeSnapshotService, Depends(service_dep)], + ) -> HomeSnapshotResponse: + try: + return await service.initialize(request) + except HomeSnapshotServiceError as exc: + raise HTTPException( + status_code=exc.status_code, + detail={"code": exc.code, "message": exc.message}, + ) from exc + + @router.post("/from-binding", response_model=HomeSnapshotResponse, status_code=status.HTTP_201_CREATED) + async def create_snapshot_from_binding( + request: CreateHomeSnapshotFromBindingRequest, + service: Annotated[HomeSnapshotService, Depends(service_dep)], + ) -> HomeSnapshotResponse: + try: + return await service.create_from_binding(request) + except HomeSnapshotServiceError as exc: + raise HTTPException( + status_code=exc.status_code, + detail={"code": exc.code, "message": exc.message}, + ) from exc + + @router.post("/delete", status_code=status.HTTP_204_NO_CONTENT) + async def delete_snapshot( + request: DeleteHomeSnapshotRequest, + service: Annotated[HomeSnapshotService, Depends(service_dep)], + ) -> Response: + try: + await service.delete(request) + except HomeSnapshotServiceError as exc: + raise HTTPException( + status_code=exc.status_code, + detail={"code": exc.code, "message": exc.message}, + ) from exc + return Response(status_code=status.HTTP_204_NO_CONTENT) + + return router + + +__all__ = ["create_home_snapshots_router"] diff --git a/dify-agent/src/dify_agent/server/routes/sandbox_files.py b/dify-agent/src/dify_agent/server/routes/sandbox_files.py deleted file mode 100644 index 10429620cc9..00000000000 --- a/dify-agent/src/dify_agent/server/routes/sandbox_files.py +++ /dev/null @@ -1,73 +0,0 @@ -"""FastAPI routes for sandbox file operations. - -The agent backend receives a structured ``SandboxLocator`` rather than a raw -shell session id. Routes stay private-network only like ``/runs`` and forward -all sandbox work to ``SandboxFileService``. -""" - -from collections.abc import Callable -from typing import Annotated - -from fastapi import APIRouter, Depends, HTTPException - -from dify_agent.protocol import ( - SandboxListRequest, - SandboxListResponse, - SandboxReadRequest, - SandboxReadResponse, - SandboxUploadRequest, - SandboxUploadResponse, -) -from dify_agent.server.sandbox_files import SandboxFileError, SandboxFileService - - -def create_sandbox_files_router(get_service: Callable[[], SandboxFileService | None]) -> APIRouter: - """Create sandbox file routes bound to the app's service provider.""" - router = APIRouter(prefix="/sandbox", tags=["sandbox"]) - - def service_dep() -> SandboxFileService: - service = get_service() - if service is None: - raise HTTPException( - status_code=503, - detail={"code": "sandbox_backend_unavailable", "message": "sandbox service is not configured"}, - ) - return service - - def raise_http(exc: SandboxFileError) -> HTTPException: - return HTTPException(status_code=exc.status_code, detail={"code": exc.code, "message": exc.message}) - - @router.post("/files/list", response_model=SandboxListResponse) - async def list_files( - request: SandboxListRequest, - service: Annotated[SandboxFileService, Depends(service_dep)], - ) -> SandboxListResponse: - try: - return await service.list_files(request) - except SandboxFileError as exc: - raise raise_http(exc) from exc - - @router.post("/files/read", response_model=SandboxReadResponse) - async def read_file( - request: SandboxReadRequest, - service: Annotated[SandboxFileService, Depends(service_dep)], - ) -> SandboxReadResponse: - try: - return await service.read_file(request) - except SandboxFileError as exc: - raise raise_http(exc) from exc - - @router.post("/files/upload", response_model=SandboxUploadResponse) - async def upload_file( - request: SandboxUploadRequest, - service: Annotated[SandboxFileService, Depends(service_dep)], - ) -> SandboxUploadResponse: - try: - return await service.upload_file(request) - except SandboxFileError as exc: - raise raise_http(exc) from exc - - return router - - -__all__ = ["create_sandbox_files_router"] diff --git a/dify-agent/src/dify_agent/server/routes/workspace_files.py b/dify-agent/src/dify_agent/server/routes/workspace_files.py new file mode 100644 index 00000000000..a1eb8d794f7 --- /dev/null +++ b/dify-agent/src/dify_agent/server/routes/workspace_files.py @@ -0,0 +1,67 @@ +"""Private Workspace file routes used by Dify API.""" + +from collections.abc import Callable +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException + +from dify_agent.protocol import ( + WorkspaceListRequest, + WorkspaceListResponse, + WorkspaceReadRequest, + WorkspaceReadResponse, + WorkspaceUploadRequest, + WorkspaceUploadResponse, +) +from dify_agent.server.workspace_files import WorkspaceFileError, WorkspaceFileService + + +def create_workspace_files_router(get_service: Callable[[], WorkspaceFileService | None]) -> APIRouter: + router = APIRouter(prefix="/workspace", tags=["workspace"]) + + def service_dep() -> WorkspaceFileService: + service = get_service() + if service is None: + raise HTTPException( + status_code=503, + detail={"code": "runtime_backend_unavailable", "message": "Workspace service is not configured"}, + ) + return service + + def raise_http(exc: WorkspaceFileError) -> HTTPException: + return HTTPException(status_code=exc.status_code, detail={"code": exc.code, "message": exc.message}) + + @router.post("/files/list", response_model=WorkspaceListResponse) + async def list_files( + request: WorkspaceListRequest, + service: Annotated[WorkspaceFileService, Depends(service_dep)], + ) -> WorkspaceListResponse: + try: + return await service.list_files(request) + except WorkspaceFileError as exc: + raise raise_http(exc) from exc + + @router.post("/files/read", response_model=WorkspaceReadResponse) + async def read_file( + request: WorkspaceReadRequest, + service: Annotated[WorkspaceFileService, Depends(service_dep)], + ) -> WorkspaceReadResponse: + try: + return await service.read_file(request) + except WorkspaceFileError as exc: + raise raise_http(exc) from exc + + @router.post("/files/upload", response_model=WorkspaceUploadResponse) + async def upload_file( + request: WorkspaceUploadRequest, + service: Annotated[WorkspaceFileService, Depends(service_dep)], + ) -> WorkspaceUploadResponse: + try: + return await service.upload_file(request) + except WorkspaceFileError as exc: + raise raise_http(exc) from exc + + return router + + +__all__ = ["create_workspace_files_router"] diff --git a/dify-agent/src/dify_agent/server/sandbox_files.py b/dify-agent/src/dify_agent/server/sandbox_files.py deleted file mode 100644 index 9d712f2db3e..00000000000 --- a/dify-agent/src/dify_agent/server/sandbox_files.py +++ /dev/null @@ -1,428 +0,0 @@ -"""Sandbox file service that re-enters prior shell sessions through the shell layer. - -Unlike the removed workspace inspector, this service never talks to shellctl -directly and never reads sandbox files outside the shell layer. It rebuilds a -minimal compositor from ``SandboxLocator``, enters the saved -``execution_context`` + ``shell`` layers, and executes fixed scripts through -``DifyShellLayer.run_remote_script_complete()``. - -The scripts still frame their structured payloads with a PTY-safe -base64-between-sentinels envelope. shellctl jobs are tmux-backed, so raw JSON can -be wrapped or surrounded by prompt noise; the framing keeps list/read/upload -responses parseable without falling back to direct shellctl file access. Path -arguments resolve from the saved shell workspace cwd, and ``~`` resolves through -the shell layer's injected sandbox ``HOME``. The scripts do not re-impose a -workspace-root boundary, so callers can use ``../`` or ``~/...`` when the sandbox -filesystem layout expects it. -""" - -from __future__ import annotations - -import json -import base64 -import binascii -import shlex -import textwrap -from dataclasses import dataclass -from typing import TypeVar, cast - -from dify_agent.layers.shell.layer import CompleteRemoteCommandResult, DifyShellLayer -from dify_agent.layers.shell.output_text import utf8_suffix -from dify_agent.protocol import ( - SandboxListRequest, - SandboxListResponse, - SandboxLocator, - SandboxReadRequest, - SandboxReadResponse, - SandboxUploadRequest, - SandboxUploadResponse, - normalize_composition, -) -from pydantic import BaseModel, ValidationError -from dify_agent.runtime.compositor_factory import DifyAgentLayerProvider, build_pydantic_ai_compositor - -_LIST_MAX_ENTRIES = 1000 -_LIST_TIMEOUT_SECONDS = 10.0 -_READ_TIMEOUT_SECONDS = 15.0 -_UPLOAD_TIMEOUT_SECONDS = 30.0 -_OUTPUT_BEGIN = "<<>>" -_OUTPUT_END = "<<>>" -_SHELL_RESULT_OUTPUT_TAIL_BYTES = 8 * 1024 -ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel) - -_LIST_SCRIPT = """ -import base64 -import json -import stat -import sys -from pathlib import Path - - -BEGIN = "<<>>" -END = "<<>>" - - -def emit(payload): - blob = base64.b64encode(json.dumps(payload, ensure_ascii=False).encode("utf-8")).decode("ascii") - print(BEGIN + blob + END) - - -raw_path = sys.argv[1] -limit = int(sys.argv[2]) -target = Path(raw_path).expanduser().resolve() - -if not target.exists(): - emit({"error": "sandbox_path_not_found", "message": "path not found in sandbox"}) - sys.exit(0) -if not target.is_dir(): - emit({"error": "sandbox_path_not_readable", "message": "path is not a directory"}) - sys.exit(0) - -entries = [] -for child in sorted(target.iterdir(), key=lambda item: item.name)[:limit]: - child_stat = child.lstat() - mode = child_stat.st_mode - if stat.S_ISLNK(mode): - entry_type = "symlink" - elif stat.S_ISDIR(mode): - entry_type = "dir" - elif stat.S_ISREG(mode): - entry_type = "file" - else: - entry_type = "other" - entries.append( - { - "name": child.name, - "type": entry_type, - "size": int(child_stat.st_size), - "mtime": int(child_stat.st_mtime), - } - ) - -emit( - { - "path": raw_path, - "entries": entries, - "truncated": len(list(target.iterdir())) > limit, - } -) -""" - -_READ_SCRIPT = """ -import base64 -import json -import sys -from pathlib import Path - - -BEGIN = "<<>>" -END = "<<>>" - - -def emit(payload): - blob = base64.b64encode(json.dumps(payload, ensure_ascii=False).encode("utf-8")).decode("ascii") - print(BEGIN + blob + END) - - -raw_path = sys.argv[1] -max_bytes = int(sys.argv[2]) -target = Path(raw_path).expanduser().resolve() -if not target.exists(): - emit({"error": "sandbox_path_not_found", "message": "path not found in sandbox"}) - sys.exit(0) -if not target.is_file(): - emit({"error": "sandbox_path_not_readable", "message": "path is not a readable file"}) - sys.exit(0) - -size = int(target.stat().st_size) -with target.open("rb") as file_obj: - data = file_obj.read(max_bytes + 1) - -truncated = len(data) > max_bytes -data = data[:max_bytes] -try: - text = data.decode("utf-8") -except UnicodeDecodeError: - emit( - { - "path": raw_path, - "size": size, - "truncated": truncated, - "binary": True, - "text": None, - } - ) - sys.exit(0) - -emit( - { - "path": raw_path, - "size": size, - "truncated": truncated, - "binary": False, - "text": text, - } -) -""" - -_UPLOAD_SCRIPT = """ -import base64 -import json -import subprocess -import sys -from pathlib import Path - - -BEGIN = "<<>>" -END = "<<>>" - - -def emit(payload): - blob = base64.b64encode(json.dumps(payload, ensure_ascii=False).encode("utf-8")).decode("ascii") - print(BEGIN + blob + END) - - -raw_path = sys.argv[1] -target = Path(raw_path).expanduser().resolve() -if not target.exists(): - emit({"error": "sandbox_path_not_found", "message": "path not found in sandbox"}) - sys.exit(0) -if not target.is_file(): - emit({"error": "sandbox_path_not_readable", "message": "path is not a readable file"}) - sys.exit(0) - -command = ["dify-agent", "file", "upload", raw_path] -completed = subprocess.run(command, capture_output=True, text=True, check=False) -if completed.returncode != 0: - emit( - { - "error": "agent_stub_upload_failed", - "message": (completed.stderr or completed.stdout or f"upload exited with code {completed.returncode}").strip(), - } - ) - sys.exit(0) - -try: - file_mapping = json.loads(completed.stdout) -except ValueError as exc: - emit({"error": "agent_stub_upload_failed", "message": f"upload returned invalid JSON: {exc}"}) - sys.exit(0) - -emit({"path": raw_path, "file": file_mapping}) -""" - - -class SandboxFileError(Exception): - """Sandbox file failure mapped to HTTP by the FastAPI route layer.""" - - code: str - message: str - status_code: int - - def __init__(self, code: str, message: str, *, status_code: int = 400) -> None: - super().__init__(message) - self.code = code - self.message = message - self.status_code = status_code - - -@dataclass(slots=True) -class SandboxFileService: - """Execute fixed sandbox file operations through the saved shell session.""" - - layer_providers: tuple[DifyAgentLayerProvider, ...] - - async def list_files(self, request: SandboxListRequest) -> SandboxListResponse: - normalized_path = _normalize_sandbox_path(request.path, allow_current_directory=True) - payload = await self._run_locator_script( - request.locator, - script_source=_LIST_SCRIPT, - args=[normalized_path, str(_LIST_MAX_ENTRIES)], - timeout=_LIST_TIMEOUT_SECONDS, - inject_agent_stub_env=False, - ) - return _validate_response_model(SandboxListResponse, payload) - - async def read_file(self, request: SandboxReadRequest) -> SandboxReadResponse: - normalized_path = _normalize_sandbox_path(request.path, allow_current_directory=False) - payload = await self._run_locator_script( - request.locator, - script_source=_READ_SCRIPT, - args=[normalized_path, str(request.max_bytes)], - timeout=_READ_TIMEOUT_SECONDS, - inject_agent_stub_env=False, - ) - return _validate_response_model(SandboxReadResponse, payload) - - async def upload_file(self, request: SandboxUploadRequest) -> SandboxUploadResponse: - normalized_path = _normalize_sandbox_path(request.path, allow_current_directory=False) - payload = await self._run_locator_script( - request.locator, - script_source=_UPLOAD_SCRIPT, - args=[normalized_path], - timeout=_UPLOAD_TIMEOUT_SECONDS, - inject_agent_stub_env=True, - ) - return _validate_response_model(SandboxUploadResponse, payload) - - async def _run_locator_script( - self, - locator: SandboxLocator, - *, - script_source: str, - args: list[str], - timeout: float, - inject_agent_stub_env: bool, - ) -> dict[str, object]: - try: - graph_config, layer_configs = normalize_composition(locator.composition) - compositor = build_pydantic_ai_compositor(graph_config, providers=self.layer_providers) - async with compositor.enter(configs=layer_configs, session_snapshot=locator.session_snapshot) as run: - run.suspend_on_exit() - shell_layer = run.get_layer("shell", DifyShellLayer) - result = await shell_layer.run_remote_script_complete( - _build_python_script_command(script_source=script_source, args=args), - timeout=timeout, - inject_agent_stub_env=inject_agent_stub_env, - ) - except (KeyError, TypeError, ValueError) as exc: - raise SandboxFileError("invalid_sandbox_locator", str(exc), status_code=400) from exc - except RuntimeError as exc: - raise SandboxFileError("sandbox_command_failed", str(exc), status_code=502) from exc - - return _decode_sandbox_payload(result) - - -def _normalize_sandbox_path(path: str, *, allow_current_directory: bool) -> str: - """Reject only syntactically unsafe paths and preserve relative traversal. - - The remote scripts run with the saved workspace cwd, so ``../`` remains a - valid sandbox-relative path when callers need to reach sibling directories. - ``~`` and ``~/...`` are also accepted and resolved by the embedded scripts - against the shell layer's sandbox ``HOME``. - """ - - normalized = (path or "").strip() - if normalized in {"", ".", "./"}: - if allow_current_directory: - return "." - raise SandboxFileError("invalid_sandbox_path", "path must not be blank", status_code=400) - if normalized.startswith("/"): - raise SandboxFileError( - "invalid_sandbox_path", "path must be relative to the sandbox workspace", status_code=400 - ) - if normalized.startswith("~") and normalized != "~" and not normalized.startswith("~/"): - raise SandboxFileError("invalid_sandbox_path", "path must use ~ or ~/ for the sandbox home", status_code=400) - if "\x00" in normalized or any(ord(ch) < 0x20 for ch in normalized): - raise SandboxFileError("invalid_sandbox_path", "path contains unsupported control characters", status_code=400) - return normalized - - -def _build_python_script_command(*, script_source: str, args: list[str]) -> str: - quoted_args = " ".join(shlex.quote(value) for value in args) - script = textwrap.dedent(script_source).strip() - return f"python3 - {quoted_args} <<'PY'\n{script}\nPY" - - -def _decode_sandbox_payload(result: CompleteRemoteCommandResult) -> dict[str, object]: - if result.exit_code not in (0, None): - raise SandboxFileError( - "sandbox_command_failed", - "sandbox command exited with code " + f"{result.exit_code}: {_shell_result_details(result)}", - status_code=502, - ) - begin = result.output.find(_OUTPUT_BEGIN) - end = result.output.find(_OUTPUT_END, begin + len(_OUTPUT_BEGIN)) if begin != -1 else -1 - if begin == -1 or end == -1: - if not result.output_complete: - raise SandboxFileError( - "sandbox_command_failed", - "sandbox command output incomplete before framed payload was captured: " - + _shell_result_details(result), - status_code=502, - ) - raise SandboxFileError( - "sandbox_command_failed", - "sandbox command returned no framed payload", - status_code=502, - ) - blob = result.output[begin + len(_OUTPUT_BEGIN) : end] - compact = "".join(blob.split()) - try: - decoded = base64.b64decode(compact, validate=True) - loaded = cast(object, json.loads(decoded.decode("utf-8"))) - except (binascii.Error, ValueError) as exc: - if not result.output_complete: - raise SandboxFileError( - "sandbox_command_failed", - "sandbox command output incomplete while decoding framed payload: " + _shell_result_details(result), - status_code=502, - ) from exc - raise SandboxFileError( - "sandbox_command_failed", - f"sandbox command returned invalid framed payload: {exc}", - status_code=502, - ) from exc - if not isinstance(loaded, dict): - if not result.output_complete: - raise SandboxFileError( - "sandbox_command_failed", - "sandbox command output incomplete while validating framed payload object: " - + _shell_result_details(result), - status_code=502, - ) - raise SandboxFileError( - "sandbox_command_failed", "sandbox command returned a non-object payload", status_code=502 - ) - payload = cast(dict[str, object], loaded) - error = payload.get("error") - if isinstance(error, str): - status_code = ( - 404 - if error in {"sandbox_not_found", "sandbox_path_not_found"} - else 502 - if error == "agent_stub_upload_failed" - else 400 - ) - if error in {"sandbox_command_failed", "agent_stub_upload_failed"}: - status_code = 502 - message = payload.get("message") - raise SandboxFileError( - error, - str(message) if isinstance(message, str) and message else error, - status_code=status_code, - ) - return payload - - -def _shell_result_details(result: CompleteRemoteCommandResult) -> str: - details = ( - f"output_complete={result.output_complete} " - + f"incomplete_reason={result.incomplete_reason} " - + f"output_path={result.output_path}" - ) - if not result.output: - return details - return details + "\n" + _bounded_output_tail(result.output) - - -def _bounded_output_tail(output: str) -> str: - tail = utf8_suffix(output, _SHELL_RESULT_OUTPUT_TAIL_BYTES) - if tail == output: - return output - return f"... (showing last {_SHELL_RESULT_OUTPUT_TAIL_BYTES} bytes of raw output) ...\n{tail}" - - -def _validate_response_model( - model_type: type[ResponseModelT], - payload: dict[str, object], -) -> ResponseModelT: - try: - return model_type.model_validate(payload) - except ValidationError as exc: - raise SandboxFileError( - "sandbox_command_failed", f"sandbox command returned invalid payload: {exc}", status_code=502 - ) from exc - - -__all__ = ["SandboxFileError", "SandboxFileService"] diff --git a/dify-agent/src/dify_agent/server/settings.py b/dify-agent/src/dify_agent/server/settings.py index c24981d0e73..40403fa7b01 100644 --- a/dify-agent/src/dify_agent/server/settings.py +++ b/dify-agent/src/dify_agent/server/settings.py @@ -12,26 +12,26 @@ live here under the ``DIFY_AGENT_...`` environment-variable namespace. import httpx -from typing import TYPE_CHECKING, ClassVar, Literal +from typing import ClassVar, Literal, cast -from pydantic import AnyHttpUrl, Field, TypeAdapter, field_validator, model_validator +from pydantic import AliasChoices, AnyHttpUrl, Field, TypeAdapter, field_validator, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict -if TYPE_CHECKING: - from dify_agent.adapters.shell.protocols import ShellProviderProtocol - from dify_agent.agent_stub.protocol.agent_stub import normalize_agent_stub_api_base_url, parse_agent_stub_endpoint from dify_agent.agent_stub.server.agent_stub_config import DifyApiAgentStubConfigRequestHandler from dify_agent.agent_stub.server.agent_stub_drive import DifyApiAgentStubDriveRequestHandler from dify_agent.agent_stub.server.agent_stub_files import DifyApiAgentStubFileRequestHandler from dify_agent.agent_stub.server.grpc_bind import normalize_agent_stub_grpc_bind_address from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec, decode_server_secret_key +from dify_agent.runtime_backend import RuntimeBackendProfile +from dify_agent.runtime_backend.e2b import E2B_MAX_ACTIVE_TIMEOUT_SECONDS +from dify_agent.runtime_backend.profile import RuntimeBackendSettings, create_runtime_backend_profile DEFAULT_RUN_RETENTION_SECONDS = 3 * 24 * 60 * 60 class ServerSettings(BaseSettings): - """Environment-backed settings for Redis, scheduling, outbound HTTP, and shell access.""" + """Environment settings for scheduling, outbound HTTP, and runtime resources.""" redis_url: str = "redis://localhost:6379/0" redis_prefix: str = "dify-agent" @@ -41,14 +41,29 @@ class ServerSettings(BaseSettings): plugin_daemon_api_key: str = "" inner_api_url: str = "http://localhost:5001" inner_api_key: str | None = None - shell_provider: Literal["shellctl", "enterprise"] = "shellctl" - shellctl_entrypoint: str | None = None - shellctl_auth_token: str | None = None - shell_home_root: str = "/home" + runtime_backend: Literal["local", "enterprise", "e2b"] = "local" + local_sandbox_endpoint: str | None = Field( + default=None, + validation_alias=AliasChoices("DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT", "DIFY_AGENT_SHELLCTL_ENTRYPOINT"), + ) + local_sandbox_auth_token: str | None = Field( + default=None, + validation_alias=AliasChoices("DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN", "DIFY_AGENT_SHELLCTL_AUTH_TOKEN"), + ) enterprise_sandbox_gateway_endpoint: str | None = None enterprise_sandbox_gateway_auth_token: str | None = None enterprise_sandbox_gateway_timeout: float = Field(default=30.0, gt=0) enterprise_sandbox_proxy_timeout: float = Field(default=60.0, gt=0) + e2b_api_key: str | None = None + e2b_template: str = "difys-default-team/dify-agent-local-sandbox" + e2b_active_timeout_seconds: int = Field( + default=E2B_MAX_ACTIVE_TIMEOUT_SECONDS, + ge=1, + le=E2B_MAX_ACTIVE_TIMEOUT_SECONDS, + ) + e2b_shellctl_auth_token: str = "" + e2b_shellctl_port: int = Field(default=5004, ge=1, le=65535) + sandbox_file_upload_max_bytes: int = Field(default=50 * 1024 * 1024, ge=1) agent_stub_api_base_url: str | None = Field(default=None, validation_alias="DIFY_AGENT_STUB_API_BASE_URL") agent_stub_grpc_bind_address: str | None = Field(default=None, validation_alias="DIFY_AGENT_STUB_GRPC_BIND_ADDRESS") server_secret_key: str | None = None @@ -135,21 +150,10 @@ class ServerSettings(BaseSettings): return [] import json as _json - parsed = _json.loads(stripped) + parsed = cast(object, _json.loads(stripped)) if not isinstance(parsed, list): raise ValueError("DIFY_AGENT_SHELL_REDACT_PATTERNS must be a JSON array of strings") - return [str(p) for p in parsed] - - @field_validator("shell_home_root") - @classmethod - def normalize_shell_home_root(cls, value: str) -> str: - """Normalize the root used for per-Agent shell HOME directories.""" - stripped = value.strip().rstrip("/") - if not stripped: - raise ValueError("DIFY_AGENT_SHELL_HOME_ROOT must not be empty") - if not stripped.startswith("/"): - raise ValueError("DIFY_AGENT_SHELL_HOME_ROOT must be an absolute path") - return stripped + return TypeAdapter(list[str]).validate_python(parsed) @model_validator(mode="after") def validate_agent_stub_requirements(self) -> "ServerSettings": @@ -165,26 +169,24 @@ class ServerSettings(BaseSettings): raise ValueError("DIFY_AGENT_STUB_GRPC_BIND_ADDRESS requires a grpc:// DIFY_AGENT_STUB_API_BASE_URL.") return self - def build_shell_provider(self) -> "ShellProviderProtocol | None": - from dify_agent.adapters.shell.config import ShellAdapterSettings - from dify_agent.adapters.shell.factory import create_shell_provider - - match self.shell_provider: - case "shellctl": - if not self.shellctl_entrypoint: - return None - case "enterprise": - if not self.enterprise_sandbox_gateway_endpoint: - return None - return create_shell_provider( - ShellAdapterSettings( - shell_provider=self.shell_provider, - shellctl_entrypoint=self.shellctl_entrypoint, - shellctl_auth_token=self.shellctl_auth_token, + def build_runtime_backend_profile(self) -> RuntimeBackendProfile | None: + """Build the deployment-selected resource backend without adding service state.""" + if self.runtime_backend == "local" and not self.local_sandbox_endpoint: + return None + return create_runtime_backend_profile( + RuntimeBackendSettings( + runtime_backend=self.runtime_backend, + local_sandbox_endpoint=self.local_sandbox_endpoint, + local_sandbox_auth_token=self.local_sandbox_auth_token, enterprise_sandbox_gateway_endpoint=self.enterprise_sandbox_gateway_endpoint, enterprise_sandbox_gateway_auth_token=self.enterprise_sandbox_gateway_auth_token, enterprise_sandbox_gateway_timeout=self.enterprise_sandbox_gateway_timeout, enterprise_sandbox_proxy_timeout=self.enterprise_sandbox_proxy_timeout, + e2b_api_key=self.e2b_api_key, + e2b_template=self.e2b_template, + e2b_active_timeout_seconds=self.e2b_active_timeout_seconds, + e2b_shellctl_auth_token=self.e2b_shellctl_auth_token, + e2b_shellctl_port=self.e2b_shellctl_port, ) ) diff --git a/dify-agent/src/dify_agent/server/workspace_files.py b/dify-agent/src/dify_agent/server/workspace_files.py new file mode 100644 index 00000000000..26e963f1092 --- /dev/null +++ b/dify-agent/src/dify_agent/server/workspace_files.py @@ -0,0 +1,222 @@ +"""Workspace file access through an operation-scoped RuntimeLease. + +Paths are passed directly to the backend FileSystem. Dify Agent does not rebase +them to ``layout.workspace_dir`` or impose another containment boundary. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +import mimetypes +from pathlib import PurePosixPath +from typing import ClassVar, Protocol + +import httpx +from pydantic import BaseModel, ConfigDict, ValidationError + +from dify_agent.agent_stub.protocol import ( + AgentStubFileDownloadRequest, + AgentStubFileMapping, + AgentStubFileUploadRequest, +) +from dify_agent.agent_stub.server.agent_stub_files import AgentStubFileRequestError, AgentStubFileRequestHandler +from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubPrincipal +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig +from dify_agent.protocol import ( + WorkspaceListRequest, + WorkspaceListResponse, + WorkspaceReadRequest, + WorkspaceReadResponse, + WorkspaceUploadRequest, + WorkspaceUploadResponse, + WorkspaceUploadedFile, +) +from dify_agent.runtime_backend import ( + BindingAcquireError, + BindingLostError, + ExecutionBindingBackend, + WorkspaceFileTooLargeError, + WorkspacePathError, + WorkspaceUnavailableError, +) +from dify_agent.runtime_backend.leases import open_runtime_lease + +_LIST_MAX_ENTRIES = 1000 +_UPLOAD_TIMEOUT_SECONDS = 30.0 + + +class _SignedUploadResponse(BaseModel): + reference: str + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class WorkspaceFileError(Exception): + code: str + message: str + status_code: int + + def __init__(self, code: str, message: str, *, status_code: int = 400) -> None: + super().__init__(message) + self.code = code + self.message = message + self.status_code = status_code + + +class WorkspaceFileUploader(Protocol): + async def upload( + self, + *, + execution_context: DifyExecutionContextLayerConfig, + filename: str, + mimetype: str, + content: bytes, + ) -> WorkspaceUploadedFile: ... + + +@dataclass(slots=True) +class AgentStubWorkspaceFileUploader: + file_request_handler: AgentStubFileRequestHandler + timeout: httpx.Timeout | float = _UPLOAD_TIMEOUT_SECONDS + transport: httpx.AsyncBaseTransport | None = None + + async def upload( + self, + *, + execution_context: DifyExecutionContextLayerConfig, + filename: str, + mimetype: str, + content: bytes, + ) -> WorkspaceUploadedFile: + principal = AgentStubPrincipal( + execution_context=execution_context, + session_id=None, + scope=[], + token_id="workspace-file-upload", + ) + try: + upload_request = await self.file_request_handler.create_upload_request( + principal=principal, + request=AgentStubFileUploadRequest(filename=filename, mimetype=mimetype), + ) + async with httpx.AsyncClient( + timeout=self.timeout, + follow_redirects=True, + trust_env=False, + transport=self.transport, + ) as client: + response = await client.post( + upload_request.upload_url, + files={"file": (filename, content, mimetype)}, + ) + _ = response.raise_for_status() + payload = _SignedUploadResponse.model_validate(response.json()) + mapping = AgentStubFileMapping(transfer_method="tool_file", reference=payload.reference) + download_request = await self.file_request_handler.create_download_request( + principal=principal, + request=AgentStubFileDownloadRequest(file=mapping, for_external=False), + ) + return WorkspaceUploadedFile( + reference=payload.reference, + download_url=download_request.download_url, + ) + except AgentStubFileRequestError as exc: + raise WorkspaceFileError("agent_stub_upload_failed", str(exc.detail), status_code=exc.status_code) from exc + except httpx.TimeoutException as exc: + raise WorkspaceFileError( + "agent_stub_upload_failed", "signed file upload timed out", status_code=504 + ) from exc + except httpx.HTTPStatusError as exc: + raise WorkspaceFileError( + "agent_stub_upload_failed", + f"signed file upload failed with status {exc.response.status_code}", + status_code=exc.response.status_code, + ) from exc + except httpx.RequestError as exc: + raise WorkspaceFileError( + "agent_stub_upload_failed", f"signed file upload failed: {exc}", status_code=502 + ) from exc + except (ValidationError, ValueError) as exc: + raise WorkspaceFileError( + "agent_stub_upload_failed", "signed file upload returned invalid data", status_code=502 + ) from exc + + +@dataclass(slots=True) +class WorkspaceFileService: + execution_bindings: ExecutionBindingBackend + upload_max_bytes: int + file_uploader: WorkspaceFileUploader | None = None + + async def list_files(self, request: WorkspaceListRequest) -> WorkspaceListResponse: + try: + async with open_runtime_lease(self.execution_bindings, request.backend_binding_ref) as lease: + result = await lease.files.list_directory(path=request.path, limit=_LIST_MAX_ENTRIES) + return WorkspaceListResponse.model_validate( + { + "path": result.path, + "entries": [asdict(entry) for entry in result.entries], + "truncated": result.truncated, + } + ) + except Exception as exc: + raise _normalize_file_error(exc) from exc + + async def read_file(self, request: WorkspaceReadRequest) -> WorkspaceReadResponse: + try: + async with open_runtime_lease(self.execution_bindings, request.backend_binding_ref) as lease: + result = await lease.files.read_file(path=request.path, max_bytes=request.max_bytes) + return WorkspaceReadResponse( + path=result.path, + size=result.size, + truncated=result.truncated, + binary=result.binary, + text=result.text, + ) + except Exception as exc: + raise _normalize_file_error(exc) from exc + + async def upload_file(self, request: WorkspaceUploadRequest) -> WorkspaceUploadResponse: + uploader = self.file_uploader + if uploader is None: + raise WorkspaceFileError( + "agent_stub_upload_unavailable", "Agent Stub file upload is not configured", status_code=503 + ) + try: + async with open_runtime_lease(self.execution_bindings, request.backend_binding_ref) as lease: + result = await lease.files.read_bytes(path=request.path, max_bytes=self.upload_max_bytes) + filename = PurePosixPath(result.path).name or "file" + mimetype = mimetypes.guess_type(filename)[0] or "application/octet-stream" + uploaded_file = await uploader.upload( + execution_context=request.execution_context, + filename=filename, + mimetype=mimetype, + content=result.content, + ) + return WorkspaceUploadResponse(path=result.path, file=uploaded_file) + except WorkspaceFileError: + raise + except Exception as exc: + raise _normalize_file_error(exc) from exc + + +def _normalize_file_error(exc: Exception) -> WorkspaceFileError: + if isinstance(exc, WorkspaceFileError): + return exc + if isinstance(exc, WorkspacePathError): + return WorkspaceFileError("invalid_workspace_path", str(exc), status_code=400) + if isinstance(exc, WorkspaceFileTooLargeError): + return WorkspaceFileError("file_too_large", str(exc), status_code=413) + if isinstance(exc, BindingLostError): + return WorkspaceFileError("binding_not_found", str(exc), status_code=404) + if isinstance(exc, (WorkspaceUnavailableError, BindingAcquireError)): + return WorkspaceFileError("workspace_unavailable", str(exc), status_code=502) + return WorkspaceFileError("workspace_file_failed", str(exc), status_code=502) + + +__all__ = [ + "AgentStubWorkspaceFileUploader", + "WorkspaceFileError", + "WorkspaceFileService", + "WorkspaceFileUploader", +] diff --git a/dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh b/dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh new file mode 100755 index 00000000000..1247636b058 --- /dev/null +++ b/dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env sh +set -eu + +script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +project_dir=$(CDPATH= cd -- "$script_dir/../../../.." && pwd) +container_name="dify-agent-runtime-backend-integration-$$" +image="${DIFY_AGENT_TEST_LOCAL_SANDBOX_IMAGE:-langgenius/dify-agent-local-sandbox:1.16.0}" +token="${DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN:-runtime-backend-integration}" + +cleanup() { + docker rm -f "$container_name" >/dev/null 2>&1 || true +} +trap cleanup EXIT INT TERM + +docker run --detach --rm \ + --name "$container_name" \ + --env SHELLCTL_ENABLE_PATH_ISOLATION=true \ + --env SHELLCTL_AUTH_TOKEN="$token" \ + --publish 127.0.0.1::5004 \ + "$image" >/dev/null + +published_address=$(docker port "$container_name" 5004/tcp | head -n 1) +endpoint="http://$published_address" +attempt=0 +until curl --fail --silent "$endpoint/healthz" >/dev/null; do + attempt=$((attempt + 1)) + if [ "$attempt" -ge 100 ]; then + docker logs "$container_name" + exit 1 + fi + sleep 0.1 +done + +cd "$project_dir" +NO_PROXY=127.0.0.1,localhost \ + DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT="$endpoint" \ + DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN="$token" \ + pdm run pytest --import-mode=importlib \ + tests/integration/dify_agent/runtime_backend/test_runtime_backend_lifecycle.py \ + -k local -q -rs "$@" diff --git a/dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py b/dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py new file mode 100644 index 00000000000..2a8927f6a4f --- /dev/null +++ b/dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py @@ -0,0 +1,221 @@ +"""Opt-in integration contracts for final Local and E2B working environments.""" + +from __future__ import annotations + +import os +import sys +import uuid + +import pytest + +from dify_agent.runtime_backend import ( + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotCreateSpec, + InitializeHomeSnapshotSpec, +) +from dify_agent.runtime_backend.e2b import E2BExecutionBindingBackend, E2BHomeSnapshotBackend, E2BSDKControlPlane +from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend + +pytestmark = pytest.mark.integration + + +def _required_env(name: str, purpose: str) -> str: + value = os.environ.get(name, "").strip() + if not value: + pytest.skip(f"set {name} to run the {purpose} integration contract") + return value + + +@pytest.mark.anyio +async def test_local_two_agents_share_workspace_but_not_home() -> None: + endpoint = _required_env("DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT", "real Local shellctl") + token = os.environ.get("DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN", "") + marker = uuid.uuid4().hex + snapshots = LocalHomeSnapshotBackend(endpoint=endpoint, auth_token=token) + bindings = LocalExecutionBindingBackend(endpoint=endpoint, auth_token=token) + snapshot_ref: str | None = None + allocations = [] + active_leases = [] + try: + snapshot_ref = await snapshots.initialize( + InitializeHomeSnapshotSpec( + tenant_id="integration-tenant", + agent_id="integration-agent", + home_snapshot_id=marker, + ) + ) + first = await bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="integration-tenant", + agent_id="agent-a", + binding_id=f"binding-a-{marker}", + workspace_id=f"workspace-{marker}", + existing_workspace_ref=None, + home_snapshot_ref=snapshot_ref, + ) + ) + allocations.append(first) + first_lease = await bindings.acquire(first.binding_ref) + active_leases.append(first_lease) + await first_lease.files.upload( + content=b"shared", remote_path="shared.txt", cwd=first_lease.layout.workspace_dir + ) + await bindings.release(first_lease) + active_leases.remove(first_lease) + + second = await bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="integration-tenant", + agent_id="agent-b", + binding_id=f"binding-b-{marker}", + workspace_id=f"workspace-{marker}", + existing_workspace_ref=first.workspace_ref, + home_snapshot_ref=snapshot_ref, + ) + ) + allocations.append(second) + second_lease = await bindings.acquire(second.binding_ref) + active_leases.append(second_lease) + shared = await second_lease.files.read_bytes(path="shared.txt", max_bytes=1024) + assert shared.content == b"shared" + assert second_lease.layout.home_dir != first_lease.layout.home_dir + assert second_lease.layout.workspace_dir == first_lease.layout.workspace_dir + await bindings.release(second_lease) + active_leases.remove(second_lease) + finally: + primary_error = sys.exc_info()[0] is not None + cleanup_errors: list[BaseException] = [] + for lease in active_leases: + try: + await bindings.release(lease) + except BaseException as exc: + cleanup_errors.append(exc) + for index, allocation in enumerate(allocations): + try: + await bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=allocation.binding_ref, + workspace_ref=allocation.workspace_ref if index == len(allocations) - 1 else None, + destroy_workspace=index == len(allocations) - 1, + ) + ) + except BaseException as exc: + cleanup_errors.append(exc) + if snapshot_ref is not None: + try: + await snapshots.delete(snapshot_ref) + except BaseException as exc: + cleanup_errors.append(exc) + if cleanup_errors and not primary_error: + raise cleanup_errors[0] + + +@pytest.mark.anyio +async def test_e2b_binding_checkpoint_and_collection() -> None: + api_key = _required_env("DIFY_AGENT_TEST_E2B_API_KEY", "real E2B") + template = os.environ.get( + "DIFY_AGENT_TEST_E2B_TEMPLATE", + "difys-default-team/dify-agent-local-sandbox", + ) + marker = uuid.uuid4().hex + control = E2BSDKControlPlane(api_key=api_key) + snapshots = E2BHomeSnapshotBackend(control_plane=control, template=template, active_timeout_seconds=3600) + bindings = E2BExecutionBindingBackend(control_plane=control, active_timeout_seconds=3600) + snapshot_ref: str | None = None + checkpoint_ref: str | None = None + allocation = None + checkpoint_allocation = None + lease = None + checkpoint_lease = None + try: + snapshot_ref = await snapshots.initialize( + InitializeHomeSnapshotSpec( + tenant_id="integration-tenant", + agent_id="integration-agent", + home_snapshot_id=marker, + ) + ) + allocation = await bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="integration-tenant", + agent_id="integration-agent", + binding_id=marker, + workspace_id=marker, + existing_workspace_ref=None, + home_snapshot_ref=snapshot_ref, + ) + ) + lease = await bindings.acquire(allocation.binding_ref) + await lease.files.upload(content=b"e2b", remote_path="probe.txt", cwd=lease.layout.workspace_dir) + await lease.files.upload(content=b"checkpoint-home", remote_path=".checkpoint-probe", cwd=lease.layout.home_dir) + assert (await lease.files.read_bytes(path="probe.txt", max_bytes=1024)).content == b"e2b" + checkpoint_ref = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec( + tenant_id="integration-tenant", + agent_id="integration-agent", + home_snapshot_id=f"checkpoint-{marker}", + ), + source=lease, + ) + await bindings.release(lease) + lease = None + + checkpoint_allocation = await bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="integration-tenant", + agent_id="integration-agent", + binding_id=f"checkpoint-{marker}", + workspace_id=f"checkpoint-{marker}", + existing_workspace_ref=None, + home_snapshot_ref=checkpoint_ref, + ) + ) + checkpoint_lease = await bindings.acquire(checkpoint_allocation.binding_ref) + restored = await checkpoint_lease.files.read_bytes(path="~/.checkpoint-probe", max_bytes=1024) + assert restored.content == b"checkpoint-home" + await bindings.release(checkpoint_lease) + checkpoint_lease = None + finally: + primary_error = sys.exc_info()[0] is not None + cleanup_errors: list[BaseException] = [] + if lease is not None: + try: + await bindings.release(lease) + except BaseException as exc: + cleanup_errors.append(exc) + if checkpoint_lease is not None: + try: + await bindings.release(checkpoint_lease) + except BaseException as exc: + cleanup_errors.append(exc) + if checkpoint_allocation is not None: + try: + await bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=checkpoint_allocation.binding_ref, + workspace_ref=checkpoint_allocation.workspace_ref, + destroy_workspace=True, + ) + ) + except BaseException as exc: + cleanup_errors.append(exc) + if allocation is not None: + try: + await bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=allocation.binding_ref, + workspace_ref=allocation.workspace_ref, + destroy_workspace=True, + ) + ) + except BaseException as exc: + cleanup_errors.append(exc) + for ref in (checkpoint_ref, snapshot_ref): + if ref is not None: + try: + await snapshots.delete(ref) + except BaseException as exc: + cleanup_errors.append(exc) + if cleanup_errors and not primary_error: + raise cleanup_errors[0] diff --git a/dify-agent/tests/local/dify_agent/adapters/shell/test_config.py b/dify-agent/tests/local/dify_agent/adapters/shell/test_config.py index 19110694403..9bbeb84f3c9 100644 --- a/dify-agent/tests/local/dify_agent/adapters/shell/test_config.py +++ b/dify-agent/tests/local/dify_agent/adapters/shell/test_config.py @@ -1,78 +1,78 @@ -"""Tests for ShellAdapterSettings provider-based field validation.""" +"""Tests for deployment-selected runtime backend settings.""" from __future__ import annotations import pytest from pydantic import ValidationError -from dify_agent.adapters.shell.config import ShellAdapterSettings +from dify_agent.runtime_backend.profile import RuntimeBackendSettings -class TestShellctlValidation: +class TestLocalValidation: def test_valid_url_passes(self) -> None: - settings = ShellAdapterSettings( - shell_provider="shellctl", - shellctl_entrypoint="http://shellctl.example.com", + settings = RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="http://shellctl.example.com", ) - assert settings.shellctl_entrypoint == "http://shellctl.example.com" + assert settings.local_sandbox_endpoint == "http://shellctl.example.com" def test_https_url_passes(self) -> None: - settings = ShellAdapterSettings( - shell_provider="shellctl", - shellctl_entrypoint="https://shellctl.internal:8443/v1", + settings = RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="https://shellctl.internal:8443/v1", ) - assert settings.shellctl_entrypoint == "https://shellctl.internal:8443/v1" + assert settings.local_sandbox_endpoint == "https://shellctl.internal:8443/v1" def test_missing_entrypoint_raises(self) -> None: - with pytest.raises(ValidationError, match="shellctl_entrypoint is required"): - ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint=None) + with pytest.raises(ValidationError, match="local_sandbox_endpoint is required"): + RuntimeBackendSettings(runtime_backend="local", local_sandbox_endpoint=None) def test_blank_entrypoint_raises(self) -> None: - with pytest.raises(ValidationError, match="shellctl_entrypoint is required"): - ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint=" ") + with pytest.raises(ValidationError, match="local_sandbox_endpoint is required"): + RuntimeBackendSettings(runtime_backend="local", local_sandbox_endpoint=" ") def test_invalid_url_raises(self) -> None: with pytest.raises(ValidationError, match="must be a valid http\\(s\\) URL"): - ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint="not-a-url") + RuntimeBackendSettings(runtime_backend="local", local_sandbox_endpoint="not-a-url") def test_ftp_scheme_rejected(self) -> None: with pytest.raises(ValidationError, match="must be a valid http\\(s\\) URL"): - ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint="ftp://host/path") + RuntimeBackendSettings(runtime_backend="local", local_sandbox_endpoint="ftp://host/path") class TestEnterpriseValidation: def test_valid_url_passes(self) -> None: - settings = ShellAdapterSettings( - shell_provider="enterprise", + settings = RuntimeBackendSettings( + runtime_backend="enterprise", enterprise_sandbox_gateway_endpoint="http://gateway.internal:9000", ) assert settings.enterprise_sandbox_gateway_endpoint == "http://gateway.internal:9000" def test_https_url_passes(self) -> None: - settings = ShellAdapterSettings( - shell_provider="enterprise", + settings = RuntimeBackendSettings( + runtime_backend="enterprise", enterprise_sandbox_gateway_endpoint="https://gateway.prod.example.com/api", ) assert settings.enterprise_sandbox_gateway_endpoint == "https://gateway.prod.example.com/api" def test_missing_endpoint_raises(self) -> None: with pytest.raises(ValidationError, match="enterprise_sandbox_gateway_endpoint is required"): - ShellAdapterSettings(shell_provider="enterprise", enterprise_sandbox_gateway_endpoint=None) + RuntimeBackendSettings(runtime_backend="enterprise", enterprise_sandbox_gateway_endpoint=None) def test_blank_endpoint_raises(self) -> None: with pytest.raises(ValidationError, match="enterprise_sandbox_gateway_endpoint is required"): - ShellAdapterSettings(shell_provider="enterprise", enterprise_sandbox_gateway_endpoint="") + RuntimeBackendSettings(runtime_backend="enterprise", enterprise_sandbox_gateway_endpoint="") def test_invalid_url_raises(self) -> None: with pytest.raises(ValidationError, match="must be a valid http\\(s\\) URL"): - ShellAdapterSettings( - shell_provider="enterprise", + RuntimeBackendSettings( + runtime_backend="enterprise", enterprise_sandbox_gateway_endpoint="just-a-hostname", ) def test_ftp_scheme_rejected(self) -> None: with pytest.raises(ValidationError, match="must be a valid http\\(s\\) URL"): - ShellAdapterSettings( - shell_provider="enterprise", + RuntimeBackendSettings( + runtime_backend="enterprise", enterprise_sandbox_gateway_endpoint="ftp://gateway/path", ) diff --git a/dify-agent/tests/local/dify_agent/adapters/shell/test_enterprise.py b/dify-agent/tests/local/dify_agent/adapters/shell/test_enterprise.py deleted file mode 100644 index 86ac7e79172..00000000000 --- a/dify-agent/tests/local/dify_agent/adapters/shell/test_enterprise.py +++ /dev/null @@ -1,42 +0,0 @@ -from unittest.mock import patch - -import httpx2 -import pytest - -from dify_agent.adapters.shell.enterprise import EnterpriseShellProvider - - -@pytest.mark.anyio -async def test_attach_uses_shellctl_compatible_http_client(monkeypatch: pytest.MonkeyPatch) -> None: - def handler(request: httpx2.Request) -> httpx2.Response: - assert request.url.path == "/proxy/v1/jobs/run" - assert request.headers["X-Sandbox-Id"] == "sandbox-1" - return httpx2.Response( - 200, - json={ - "job_id": "job-1", - "done": True, - "status": "exited", - "exit_code": 0, - "output_path": "/tmp/output.log", - "output": "", - "offset": 0, - "truncated": False, - }, - ) - - transport = httpx2.MockTransport(handler) - monkeypatch.setattr(httpx2, "AsyncHTTPTransport", lambda **kwargs: transport) - with patch.object(httpx2, "AsyncClient", wraps=httpx2.AsyncClient) as async_client: - provider = EnterpriseShellProvider( - gateway_endpoint="http://gateway.example", - auth_token="secret", - proxy_timeout=90, - ) - - resource = await provider.attach("sandbox-1") - - timeout = async_client.call_args.kwargs["timeout"] - assert isinstance(timeout, httpx2.Timeout) - assert timeout.read == 90 - await resource.suspend() diff --git a/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py b/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py index 9af8a15dbb1..44f897a17c0 100644 --- a/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py +++ b/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py @@ -6,21 +6,29 @@ import asyncio import base64 from collections.abc import Callable from dataclasses import dataclass, field +import json +import os +from pathlib import Path +import signal +import subprocess +import sys from typing import cast import httpx2 as httpx import pytest -from pydantic import ValidationError +from shellctl.client import ShellctlClientError from dify_agent.adapters.shell import shellctl -from dify_agent.adapters.shell.config import ShellAdapterSettings -from dify_agent.adapters.shell.factory import create_shell_provider from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellProviderError from dify_agent.adapters.shell.shellctl import ( ShellctlClientProtocol, + ShellctlCommands, ShellFileTransferError, - ShellctlProvider, + ShellctlFileTransfer, ) +from dify_agent.runtime_backend.errors import WorkspaceFileTooLargeError + +_WORKSPACE_SCRIPT_TIMEOUT_SECONDS = 5.0 @dataclass(slots=True) @@ -129,46 +137,199 @@ class FakeShellctlClient: self.closed = True -def _provider(client: FakeShellctlClient) -> ShellctlProvider: - return ShellctlProvider( - entrypoint="http://shellctl", - token="", - client_factory=lambda: _client_protocol(client), - ) - - def _client_protocol(client: FakeShellctlClient) -> ShellctlClientProtocol: return cast(ShellctlClientProtocol, cast(object, client)) -def test_factory_unknown_provider_raises() -> None: - with pytest.raises(ValidationError): - ShellAdapterSettings(shell_provider="nope") # type: ignore[arg-type] +def _kill_process_group(process: subprocess.Popen[str]) -> None: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass -def test_factory_shellctl_requires_entrypoint() -> None: - with pytest.raises(ValidationError, match="shellctl_entrypoint is required"): - ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint=None) +def _run_workspace_source(source: str, args: list[str]) -> dict[str, object]: + process = subprocess.Popen( + [sys.executable, "-c", source, *args], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = process.communicate(timeout=_WORKSPACE_SCRIPT_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + _kill_process_group(process) + _ = process.communicate() + pytest.fail(f"workspace script exceeded {_WORKSPACE_SCRIPT_TIMEOUT_SECONDS:g}s test timeout") + except BaseException: + _kill_process_group(process) + _ = process.communicate() + raise + assert process.returncode == 0, stderr + begin = stdout.find(shellctl._WORKSPACE_PAYLOAD_BEGIN) + end = stdout.find(shellctl._WORKSPACE_PAYLOAD_END, begin + len(shellctl._WORKSPACE_PAYLOAD_BEGIN)) + assert begin >= 0 and end >= 0 + encoded = "".join(stdout[begin + len(shellctl._WORKSPACE_PAYLOAD_BEGIN) : end].split()) + payload = json.loads(base64.b64decode(encoded, validate=True)) + assert isinstance(payload, dict) + return cast(dict[str, object], payload) -def test_factory_builds_shellctl_provider_from_settings() -> None: - settings = ShellAdapterSettings(shell_provider="shellctl", shellctl_entrypoint="http://shellctl.example") - provider = create_shell_provider(settings) - assert isinstance(provider, ShellctlProvider) - assert provider.entrypoint == "http://shellctl.example" - assert provider.token == "" +def _inject_workspace_checkpoint(source: str, checkpoint: str, injected_source: str) -> str: + marker = f"# DIFY_WORKSPACE_CHECKPOINT: {checkpoint}" + assert source.count(marker) == 1 + marker_index = source.index(marker) + line_start = source.rfind("\n", 0, marker_index) + 1 + indentation = source[line_start:marker_index] + assert not indentation.strip() + target = f"{indentation}{marker}" + injected = "\n".join(f"{indentation}{line}" if line else "" for line in injected_source.splitlines()) + instrumented = source.replace(target, f"{injected}\n{target}", 1) + assert instrumented != source + assert instrumented.count(marker) == 1 + return instrumented -def test_provider_create_opens_only_live_resource_and_suspend_closes_client() -> None: - client = FakeShellctlClient() +def test_workspace_read_holds_open_fd_across_concurrent_symlink_swap(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + reports = workspace / "reports" + reports.mkdir(parents=True) + source = reports / "result.pdf" + _ = source.write_text("workspace-content") + outside = tmp_path / "outside.pdf" + _ = outside.write_text("outside-content") + instrumented = _inject_workspace_checkpoint( + shellctl._READ_WORKSPACE_SCRIPT, + "file_opened", + "os.unlink(sys.argv[3])\nos.symlink(sys.argv[4], sys.argv[3])", + ) - async def scenario() -> None: - resource = await _provider(client).create() - assert client.run_calls == [] - await resource.suspend() + payload = _run_workspace_source( + instrumented, + [str(source), "1024", str(source), str(outside)], + ) - asyncio.run(scenario()) - assert client.closed is True + assert payload["text"] == "workspace-content" + assert source.is_symlink() + + +def test_workspace_read_bytes_holds_open_fd_across_same_uid_symlink_swap(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + reports = workspace / "reports" + reports.mkdir(parents=True) + source = reports / "result.pdf" + source.write_bytes(b"workspace-content\x00") + outside = tmp_path / "outside.pdf" + outside.write_bytes(b"outside-content") + instrumented = _inject_workspace_checkpoint( + shellctl._READ_WORKSPACE_BYTES_SCRIPT, + "file_opened", + "os.unlink(sys.argv[3])\nos.symlink(sys.argv[4], sys.argv[3])", + ) + + payload = _run_workspace_source( + instrumented, + [str(source), "1024", str(source), str(outside)], + ) + + encoded = payload["content_base64"] + assert isinstance(encoded, str) + assert base64.b64decode(encoded) == b"workspace-content\x00" + assert source.is_symlink() + + +def test_workspace_read_bytes_accepts_file_at_size_limit(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + source = workspace / "result.bin" + source.write_bytes(b"12345") + + payload = _run_workspace_source( + shellctl._READ_WORKSPACE_BYTES_SCRIPT, + [str(source), "5"], + ) + + encoded = payload["content_base64"] + assert isinstance(encoded, str) + assert base64.b64decode(encoded) == b"12345" + assert payload["size"] == 5 + + +def test_workspace_read_bytes_rejects_oversize_before_remote_read(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + source = workspace / "result.bin" + source.write_bytes(b"123456") + instrumented = _inject_workspace_checkpoint( + shellctl._READ_WORKSPACE_BYTES_SCRIPT, + "arguments_loaded", + "os.read = lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError('must not read'))\n\n", + ) + + payload = _run_workspace_source( + instrumented, + [str(source), "5"], + ) + + assert payload == {"path": str(source), "size": 6, "too_large": True} + + +def test_workspace_read_bytes_caps_capture_when_file_grows_after_fstat(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + source = workspace / "result.bin" + source.write_bytes(b"12345") + instrumented = _inject_workspace_checkpoint( + shellctl._READ_WORKSPACE_BYTES_SCRIPT, + "arguments_loaded", + "real_read = os.read\n" + "captured_bytes = 0\n" + "def bounded_read(fd, count):\n" + " global captured_bytes\n" + " data = real_read(fd, count)\n" + " captured_bytes += len(data)\n" + " assert captured_bytes <= int(sys.argv[2]) + 1\n" + " return data\n" + "os.read = bounded_read", + ) + instrumented = _inject_workspace_checkpoint( + instrumented, + "file_size_captured", + "with open(sys.argv[3], 'ab') as growing:\n growing.write(b'x' * 10000)", + ) + + payload = _run_workspace_source( + instrumented, + [str(source), "5", str(source)], + ) + + assert payload == {"path": str(source), "size": 6, "too_large": True} + + +def test_workspace_list_holds_open_directory_fd_across_concurrent_symlink_swap(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + reports = workspace / "reports" + reports.mkdir(parents=True) + _ = (reports / "safe.txt").write_text("safe") + outside = tmp_path / "outside" + outside.mkdir() + _ = (outside / "secret.txt").write_text("secret") + moved = workspace / "opened-reports" + instrumented = _inject_workspace_checkpoint( + shellctl._LIST_WORKSPACE_SCRIPT, + "directory_opened", + "os.rename(sys.argv[3], sys.argv[4])\nos.symlink(sys.argv[5], sys.argv[3])", + ) + + payload = _run_workspace_source( + instrumented, + [str(reports), "100", str(reports), str(moved), str(outside)], + ) + + entries = payload["entries"] + assert isinstance(entries, list) + assert [entry["name"] for entry in entries if isinstance(entry, dict)] == ["safe.txt"] def test_commands_forward_parameters_and_map_metadata() -> None: @@ -223,15 +384,14 @@ def test_commands_forward_parameters_and_map_metadata() -> None: ) async def scenario() -> None: - resource = await _provider(client).create() - run_result = await resource.commands.run("pwd", cwd="~/workspace/abc12ff", env={"FOO": "bar"}, timeout=2.5) - wait_result = await resource.commands.wait("run-job", offset=3, timeout=4.0) - read_result = await resource.commands.read_output("run-job", offset=6) - input_result = await resource.commands.input("run-job", "ls\n", offset=6, timeout=5.0) - interrupt_result = await resource.commands.interrupt("run-job", grace_seconds=1.5) - tail_result = await resource.commands.tail("run-job") - await resource.commands.delete("run-job", force=True, grace_seconds=2.0) - await resource.suspend() + commands = ShellctlCommands(_client_protocol(client)) + run_result = await commands.run("pwd", cwd="~/workspace/abc12ff", env={"FOO": "bar"}, timeout=2.5) + wait_result = await commands.wait("run-job", offset=3, timeout=4.0) + read_result = await commands.read_output("run-job", offset=6) + input_result = await commands.input("run-job", "ls\n", offset=6, timeout=5.0) + interrupt_result = await commands.interrupt("run-job", grace_seconds=1.5) + tail_result = await commands.tail("run-job") + await commands.delete("run-job", force=True, grace_seconds=2.0) assert run_result == ShellCommandResult( job_id="run-job", @@ -261,6 +421,38 @@ def test_commands_forward_parameters_and_map_metadata() -> None: assert client.delete_calls == [("run-job", True, 2.0)] +def test_commands_enforce_runtime_lease_home_and_cwd_namespace() -> None: + client = FakeShellctlClient() + + async def scenario() -> None: + commands = ShellctlCommands( + _client_protocol(client), + home_dir="/homes/binding-b", + workspace_dir="/workspaces/shared", + ) + await commands.run("pwd", env={"HOME": "/homes/binding-a", "FOO": "bar"}, timeout=2.5) + await commands.run("pwd", cwd="~/project", timeout=2.5) + with pytest.raises(ValueError, match="outside this RuntimeLease"): + await commands.run("cat secret", cwd="/homes/binding-a", timeout=2.5) + + asyncio.run(scenario()) + + assert client.run_calls == [ + _RunCall( + script="pwd", + cwd="/workspaces/shared", + env={"HOME": "/homes/binding-b", "FOO": "bar"}, + timeout=2.5, + ), + _RunCall( + script="pwd", + cwd="/homes/binding-b/project", + env={"HOME": "/homes/binding-b"}, + timeout=2.5, + ), + ] + + def test_commands_map_http_timeout_to_shell_provider_error() -> None: request = httpx.Request("POST", "http://shellctl.example/v1/jobs") client = FakeShellctlClient( @@ -270,9 +462,9 @@ def test_commands_map_http_timeout_to_shell_provider_error() -> None: ) async def scenario() -> None: - resource = await _provider(client).create() + commands = ShellctlCommands(_client_protocol(client)) with pytest.raises(ShellProviderError, match="timed out") as exc_info: - await resource.commands.run("pwd", timeout=2.5) + await commands.run("pwd", timeout=2.5) assert exc_info.value.code == "timeout" asyncio.run(scenario()) @@ -287,14 +479,84 @@ def test_commands_map_http_request_error_to_shell_provider_error() -> None: ) async def scenario() -> None: - resource = await _provider(client).create() + commands = ShellctlCommands(_client_protocol(client)) with pytest.raises(ShellProviderError, match="connection failed") as exc_info: - await resource.commands.wait("run-job", offset=3, timeout=4.0) + await commands.wait("run-job", offset=3, timeout=4.0) assert exc_info.value.code == "request_error" asyncio.run(scenario()) +def test_commands_preserve_shellctl_structured_error_fields() -> None: + client = FakeShellctlClient( + run_handler=lambda script, cwd, env, timeout: (_ for _ in ()).throw( + ShellctlClientError(404, "sandbox_not_found", "sandbox expired") + ) + ) + + async def scenario() -> None: + commands = ShellctlCommands(_client_protocol(client)) + with pytest.raises(ShellProviderError, match="sandbox expired") as exc_info: + await commands.run("pwd", timeout=2.5) + assert exc_info.value.status_code == 404 + assert exc_info.value.code == "sandbox_not_found" + + asyncio.run(scenario()) + + +def test_read_bytes_maps_oversize_payload_to_domain_error() -> None: + payload = base64.b64encode(b'{"path":"reports/large.bin","size":6,"too_large":true}').decode("ascii") + client = FakeShellctlClient( + run_handler=lambda script, cwd, env, timeout: _Job( + job_id="read-job", + status="exited", + done=True, + exit_code=0, + output=f"{shellctl._WORKSPACE_PAYLOAD_BEGIN}{payload}{shellctl._WORKSPACE_PAYLOAD_END}", + ) + ) + + async def scenario() -> None: + files = ShellctlFileTransfer(_client_protocol(client), cwd="/workspace", home_dir="/home/dify") + with pytest.raises(WorkspaceFileTooLargeError) as exc_info: + await files.read_bytes(path="reports/large.bin", max_bytes=5) + assert exc_info.value.size == 6 + assert exc_info.value.max_bytes == 5 + + asyncio.run(scenario()) + + +def test_file_operations_use_runtime_lease_namespace() -> None: + payload = base64.b64encode(b'{"path":"/homes/binding-b/report.txt","size":2,"content_base64":"b2s="}').decode( + "ascii" + ) + client = FakeShellctlClient( + run_handler=lambda script, cwd, env, timeout: _Job( + job_id="read-job", + status="exited", + done=True, + exit_code=0, + output=f"{shellctl._WORKSPACE_PAYLOAD_BEGIN}{payload}{shellctl._WORKSPACE_PAYLOAD_END}", + ) + ) + + async def scenario() -> None: + files = ShellctlFileTransfer( + _client_protocol(client), + cwd="/workspaces/shared", + home_dir="/homes/binding-b", + ) + result = await files.read_bytes(path="~/report.txt", max_bytes=10) + assert result.content == b"ok" + with pytest.raises(ValueError, match="outside this RuntimeLease"): + await files.download(remote_path="secret", cwd="/homes/binding-a") + + asyncio.run(scenario()) + + assert client.run_calls[0].cwd == "/workspaces/shared" + assert client.run_calls[0].env == {"HOME": "/homes/binding-b"} + + def test_delete_maps_http_timeout_to_shell_provider_error() -> None: request = httpx.Request("DELETE", "http://shellctl.example/v1/jobs/run-job") @@ -307,9 +569,9 @@ def test_delete_maps_http_timeout_to_shell_provider_error() -> None: client = DeleteTimeoutClient() async def scenario() -> None: - resource = await _provider(client).create() + commands = ShellctlCommands(_client_protocol(client)) with pytest.raises(ShellProviderError, match="delete timed out") as exc_info: - await resource.commands.delete("run-job", force=True, grace_seconds=2.0) + await commands.delete("run-job", force=True, grace_seconds=2.0) assert exc_info.value.code == "timeout" asyncio.run(scenario()) @@ -328,9 +590,9 @@ def test_delete_maps_http_request_error_to_shell_provider_error() -> None: client = DeleteRequestErrorClient() async def scenario() -> None: - resource = await _provider(client).create() + commands = ShellctlCommands(_client_protocol(client)) with pytest.raises(ShellProviderError, match="delete connection failed") as exc_info: - await resource.commands.delete("run-job", force=True, grace_seconds=2.0) + await commands.delete("run-job", force=True, grace_seconds=2.0) assert exc_info.value.code == "request_error" asyncio.run(scenario()) @@ -355,9 +617,9 @@ def test_files_upload_and_download_still_work() -> None: ) async def scenario() -> None: - resource = await _provider(client).create() - await resource.files.upload(content=content, remote_path="out.bin", cwd="~/workspace/abc12ff") - downloaded = await resource.files.download(remote_path="report.txt", cwd="~/workspace/abc12ff") + files = ShellctlFileTransfer(_client_protocol(client)) + await files.upload(content=content, remote_path="out.bin", cwd="~/workspace/abc12ff") + downloaded = await files.download(remote_path="report.txt", cwd="~/workspace/abc12ff") assert downloaded == content asyncio.run(scenario()) @@ -439,8 +701,8 @@ def test_download_missing_file_raises() -> None: ) async def scenario() -> None: - resource = await _provider(client).create() + files = ShellctlFileTransfer(_client_protocol(client)) with pytest.raises(ShellFileTransferError, match="not found"): - await resource.files.download(remote_path="missing.txt") + await files.download(remote_path="missing.txt") asyncio.run(scenario()) diff --git a/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py b/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py index 54e7863aeb4..8782ba48af0 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py @@ -75,6 +75,26 @@ def test_dify_api_agent_stub_file_handler_injects_execution_context_for_upload(m asyncio.run(scenario()) +def test_dify_api_agent_stub_file_handler_resolves_relative_upload_url(monkeypatch) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": {"url": "/files/upload/for-plugin?signed=yes"}}) + + _patch_async_client(monkeypatch, handler) + file_handler = DifyApiAgentStubFileRequestHandler( + inner_api_url="https://api.example.com", + inner_api_key="inner-secret", + ) + + async def scenario() -> None: + response = await file_handler.create_upload_request( + principal=_principal(), + request=AgentStubFileUploadRequest(filename="report.pdf", mimetype="application/pdf"), + ) + assert response.upload_url == "https://api.example.com/files/upload/for-plugin?signed=yes" + + asyncio.run(scenario()) + + def test_dify_api_agent_stub_file_handler_injects_execution_context_for_download(monkeypatch) -> None: def handler(request: httpx.Request) -> httpx.Response: assert str(request.url) == "https://api.example.com/inner/api/download/file/request" diff --git a/dify-agent/tests/local/dify_agent/client/test_client.py b/dify-agent/tests/local/dify_agent/client/test_client.py index 37d5c540987..97811b8ea47 100644 --- a/dify-agent/tests/local/dify_agent/client/test_client.py +++ b/dify-agent/tests/local/dify_agent/client/test_client.py @@ -11,6 +11,7 @@ import pytest from agenton.compositor import CompositorSessionSnapshot from agenton_collections.layers.plain import PLAIN_PROMPT_LAYER_TYPE_ID +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.client import _client as client_module from dify_agent.client import ( Client, @@ -23,7 +24,11 @@ from dify_agent.client import ( from dify_agent.protocol import ( CancelRunRequest, CancelRunResponse, + CreateExecutionBindingRequest, + CreateHomeSnapshotFromBindingRequest, CreateRunRequest, + DestroyExecutionBindingRequest, + InitializeHomeSnapshotRequest, RUN_EVENT_ADAPTER, RunCancelledEvent, RunEvent, @@ -33,10 +38,10 @@ from dify_agent.protocol import ( RunStartedEvent, RunSucceededEvent, RunSucceededEventData, - SandboxListResponse, - SandboxLocator, - SandboxReadResponse, - SandboxUploadResponse, + WorkspaceListResponse, + WorkspaceReadResponse, + WorkspaceUploadRequest, + WorkspaceUploadResponse, ) @@ -79,41 +84,16 @@ def _run_status_json(status: str) -> dict[str, object]: return {"run_id": "run-1", "status": status, "created_at": now, "updated_at": now, "error": None} -def _sandbox_locator() -> SandboxLocator: - return SandboxLocator.model_validate( - { - "composition": { - "schema_version": 1, - "layers": [ - { - "name": "execution_context", - "type": "dify.execution_context", - "config": { - "tenant_id": "tenant-1", - "user_from": "account", - "agent_mode": "agent_app", - "invoke_from": "service-api", - }, - }, - { - "name": "shell", - "type": "dify.shell", - "deps": {"execution_context": "execution_context"}, - "config": {}, - }, - ], - }, - "session_snapshot": { - "layers": [ - {"name": "execution_context", "lifecycle_state": "suspended", "runtime_state": {}}, - { - "name": "shell", - "lifecycle_state": "suspended", - "runtime_state": {"session_id": "abc12ff", "workspace_cwd": "~/workspace/abc12ff"}, - }, - ] - }, - } +def _workspace_upload_request(path: str = "report.txt") -> WorkspaceUploadRequest: + return WorkspaceUploadRequest( + backend_binding_ref="binding-ref", + path=path, + execution_context=DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_from="account", + agent_mode="agent_app", + invoke_from="debugger", + ), ) @@ -247,22 +227,21 @@ def test_async_methods_and_wait_run_parse_protocol_dtos() -> None: asyncio.run(scenario()) -def test_sync_sandbox_methods_post_dtos_and_parse_responses() -> None: - locator = _sandbox_locator() - +def test_sync_workspace_methods_post_dtos_and_parse_responses() -> None: def handler(request: httpx.Request) -> httpx.Response: - if request.url.path == "/sandbox/files/list": + if request.url.path == "/workspace/files/list": payload = cast(dict[str, object], json.loads(request.content)) assert payload["path"] == "." + assert payload["backend_binding_ref"] == "binding-ref" return httpx.Response(200, json={"path": ".", "entries": [], "truncated": False}) - if request.url.path == "/sandbox/files/read": + if request.url.path == "/workspace/files/read": payload = cast(dict[str, object], json.loads(request.content)) assert payload["path"] == "note.txt" assert payload["max_bytes"] == 128 return httpx.Response( 200, json={"path": "note.txt", "size": 5, "truncated": False, "binary": False, "text": "hello"} ) - if request.url.path == "/sandbox/files/upload": + if request.url.path == "/workspace/files/upload": payload = cast(dict[str, object], json.loads(request.content)) assert payload["path"] == "report.txt" return httpx.Response( @@ -280,30 +259,28 @@ def test_sync_sandbox_methods_post_dtos_and_parse_responses() -> None: client = Client(base_url="http://testserver", sync_http_client=httpx.Client(transport=httpx.MockTransport(handler))) - listing = client.list_sandbox_files_sync(locator, ".") - preview = client.read_sandbox_file_sync(locator, "note.txt", max_bytes=128) - uploaded = client.upload_sandbox_file_sync(locator, "report.txt") + listing = client.list_workspace_files_sync("binding-ref", ".") + preview = client.read_workspace_file_sync("binding-ref", "note.txt", max_bytes=128) + uploaded = client.upload_workspace_file_sync(_workspace_upload_request()) - assert isinstance(listing, SandboxListResponse) + assert isinstance(listing, WorkspaceListResponse) assert listing.path == "." - assert isinstance(preview, SandboxReadResponse) + assert isinstance(preview, WorkspaceReadResponse) assert preview.text == "hello" - assert isinstance(uploaded, SandboxUploadResponse) + assert isinstance(uploaded, WorkspaceUploadResponse) assert uploaded.file.reference == "dify-file-ref:file-1" assert uploaded.file.download_url == "https://files.example.com/report.txt" -def test_async_sandbox_methods_post_dtos_and_parse_responses() -> None: - locator = _sandbox_locator() - +def test_async_workspace_methods_post_dtos_and_parse_responses() -> None: def handler(request: httpx.Request) -> httpx.Response: - if request.url.path == "/sandbox/files/list": + if request.url.path == "/workspace/files/list": return httpx.Response(200, json={"path": ".", "entries": [], "truncated": False}) - if request.url.path == "/sandbox/files/read": + if request.url.path == "/workspace/files/read": return httpx.Response( 200, json={"path": "note.txt", "size": 5, "truncated": False, "binary": False, "text": "hello"} ) - if request.url.path == "/sandbox/files/upload": + if request.url.path == "/workspace/files/upload": return httpx.Response( 200, json={ @@ -321,9 +298,9 @@ def test_async_sandbox_methods_post_dtos_and_parse_responses() -> None: http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) client = Client(base_url="http://testserver", async_http_client=http_client) - listing = await client.list_sandbox_files(locator, ".") - preview = await client.read_sandbox_file(locator, "note.txt") - uploaded = await client.upload_sandbox_file(locator, "report.txt") + listing = await client.list_workspace_files("binding-ref", ".") + preview = await client.read_workspace_file("binding-ref", "note.txt") + uploaded = await client.upload_workspace_file(_workspace_upload_request()) assert listing.path == "." assert preview.text == "hello" @@ -334,11 +311,148 @@ def test_async_sandbox_methods_post_dtos_and_parse_responses() -> None: asyncio.run(scenario()) -def test_sync_upload_sandbox_file_rejects_missing_download_url() -> None: - locator = _sandbox_locator() +def _initialize_home_snapshot_request() -> InitializeHomeSnapshotRequest: + return InitializeHomeSnapshotRequest( + tenant_id="tenant-1", + agent_id="agent-1", + home_snapshot_id="home-1", + ) + +def test_sync_execution_binding_client_uses_private_binding_routes() -> None: def handler(request: httpx.Request) -> httpx.Response: - if request.url.path != "/sandbox/files/upload": + payload = cast(dict[str, object], json.loads(request.content)) + if request.url.path == "/execution-bindings": + assert payload["binding_id"] == "binding-1" + assert payload["home_snapshot_ref"] == "home-ref" + return httpx.Response(201, json={"binding_ref": "backend-binding", "workspace_ref": "backend-workspace"}) + assert request.url.path == "/execution-bindings/destroy" + assert payload == { + "binding_ref": "backend-binding", + "destroy_workspace": True, + "workspace_ref": "backend-workspace", + } + return httpx.Response(204) + + client = Client( + base_url="http://testserver", + sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + allocation = client.create_execution_binding_sync( + CreateExecutionBindingRequest( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="home-ref", + ) + ) + client.destroy_execution_binding_sync( + DestroyExecutionBindingRequest( + binding_ref=allocation.binding_ref, + workspace_ref=allocation.workspace_ref, + destroy_workspace=True, + ) + ) + + assert allocation.binding_ref == "backend-binding" + + +def _create_home_snapshot_from_binding_request() -> CreateHomeSnapshotFromBindingRequest: + return CreateHomeSnapshotFromBindingRequest( + tenant_id="tenant-1", + agent_id="agent-1", + home_snapshot_id="home-2", + backend_binding_ref="binding-ref", + ) + + +def test_sync_home_snapshot_client_parses_initialize_checkpoint_and_delete() -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + if request.url.path == "/home-snapshots/initialize": + assert json.loads(request.content) == _initialize_home_snapshot_request().model_dump(mode="json") + return httpx.Response(201, json={"snapshot_ref": "initial-home"}) + if request.url.path == "/home-snapshots/from-binding": + assert json.loads(request.content) == _create_home_snapshot_from_binding_request().model_dump( + mode="json" + ) + return httpx.Response(201, json={"snapshot_ref": "team/home 1"}) + assert request.url.path == "/home-snapshots/delete" + assert json.loads(request.content) == {"snapshot_ref": "team/home 1"} + return httpx.Response(204) + raise AssertionError(request.url) + + http_client = httpx.Client(transport=httpx.MockTransport(handler)) + client = Client(base_url="http://testserver", sync_http_client=http_client) + + initialized = client.initialize_home_snapshot_sync(_initialize_home_snapshot_request()) + created = client.create_home_snapshot_from_binding_sync(_create_home_snapshot_from_binding_request()) + client.delete_home_snapshot_sync(created.snapshot_ref) + + assert initialized.snapshot_ref == "initial-home" + assert created.snapshot_ref == "team/home 1" + http_client.close() + + +def test_async_home_snapshot_client_parses_initialize_checkpoint_and_delete() -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + if request.url.path == "/home-snapshots/initialize": + return httpx.Response(201, json={"snapshot_ref": "initial-home"}) + if request.url.path == "/home-snapshots/from-binding": + return httpx.Response(201, json={"snapshot_ref": "team/home 1"}) + assert request.url.path == "/home-snapshots/delete" + return httpx.Response(204) + raise AssertionError(request.url) + + async def scenario() -> None: + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + client = Client(base_url="http://testserver", async_http_client=http_client) + + initialized = await client.initialize_home_snapshot(_initialize_home_snapshot_request()) + created = await client.create_home_snapshot_from_binding(_create_home_snapshot_from_binding_request()) + await client.delete_home_snapshot(created.snapshot_ref) + + assert initialized.snapshot_ref == "initial-home" + assert created.snapshot_ref == "team/home 1" + await http_client.aclose() + + asyncio.run(scenario()) + + +def test_home_snapshot_client_maps_sync_validation_and_async_http_errors() -> None: + sync_http_client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json={}))) + sync_client = Client( + base_url="http://testserver", + sync_http_client=sync_http_client, + ) + + with pytest.raises(DifyAgentValidationError): + _ = sync_client.initialize_home_snapshot_sync(_initialize_home_snapshot_request()) + sync_http_client.close() + + async def scenario() -> None: + http_client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda _request: httpx.Response(502, json={"detail": {"code": "backend_failed"}}) + ) + ) + client = Client(base_url="http://testserver", async_http_client=http_client) + + with pytest.raises(DifyAgentHTTPError) as exc_info: + _ = await client.create_home_snapshot_from_binding(_create_home_snapshot_from_binding_request()) + assert exc_info.value.status_code == 502 + assert exc_info.value.detail == {"code": "backend_failed"} + await http_client.aclose() + + asyncio.run(scenario()) + + +def test_sync_upload_workspace_file_rejects_missing_download_url() -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path != "/workspace/files/upload": raise AssertionError(f"unexpected request: {request.method} {request.url}") return httpx.Response( 200, @@ -354,10 +468,10 @@ def test_sync_upload_sandbox_file_rejects_missing_download_url() -> None: client = Client(base_url="http://testserver", sync_http_client=httpx.Client(transport=httpx.MockTransport(handler))) with pytest.raises(DifyAgentValidationError): - _ = client.upload_sandbox_file_sync(locator, "report.txt") + _ = client.upload_workspace_file_sync(_workspace_upload_request()) -def test_sync_sandbox_methods_map_invalid_json_to_validation_error() -> None: +def test_sync_workspace_methods_map_invalid_json_to_validation_error() -> None: responses = iter([httpx.Response(200, text="not-json"), httpx.Response(404, json={"detail": "missing"})]) def handler(_request: httpx.Request) -> httpx.Response: @@ -366,10 +480,10 @@ def test_sync_sandbox_methods_map_invalid_json_to_validation_error() -> None: client = Client(base_url="http://testserver", sync_http_client=httpx.Client(transport=httpx.MockTransport(handler))) with pytest.raises(DifyAgentValidationError): - _ = client.list_sandbox_files_sync(_sandbox_locator(), ".") + _ = client.list_workspace_files_sync("binding-ref", ".") with pytest.raises(DifyAgentHTTPError) as http_error: - _ = client.read_sandbox_file_sync(_sandbox_locator(), "missing.txt") + _ = client.read_workspace_file_sync("binding-ref", "missing.txt") assert http_error.value.status_code == 404 diff --git a/dify-agent/tests/local/dify_agent/layers/config/test_layer.py b/dify-agent/tests/local/dify_agent/layers/config/test_layer.py index f085f6cc51b..c222ddaab3e 100644 --- a/dify-agent/tests/local/dify_agent/layers/config/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/config/test_layer.py @@ -7,7 +7,6 @@ from typing import Literal import pytest -from dify_agent.adapters.shell.shellctl import ShellctlProvider from dify_agent.layers.config import DifyConfigLayerConfig from dify_agent.layers.config.layer import ( DifyConfigLayer, @@ -18,14 +17,9 @@ from dify_agent.layers.shell import DifyShellLayerConfig from dify_agent.layers.shell.layer import CompleteRemoteCommandResult, DifyShellLayer -def _unused_client_factory(): - raise AssertionError("shellctl client should not be used by these config-layer tests") - - def _shell_layer() -> DifyShellLayer: return DifyShellLayer.from_config_with_settings( DifyShellLayerConfig(agent_stub_drive_ref="agent-1"), - shell_provider=ShellctlProvider(entrypoint="http://shellctl", token="", client_factory=_unused_client_factory), ) diff --git a/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py b/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py index afcef5cc396..c4cfb74346f 100644 --- a/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py @@ -6,21 +6,15 @@ from typing import Literal import pytest -from dify_agent.adapters.shell.shellctl import ShellctlProvider from dify_agent.layers.drive import DifyDriveLayerConfig, DifyDriveSkillConfig from dify_agent.layers.drive.layer import DifyDriveLayer, DifyDriveLayerError, _AGENT_FILE_UPLOAD_REPLY_HINT from dify_agent.layers.shell import DifyShellLayerConfig from dify_agent.layers.shell.layer import CompleteRemoteCommandResult, DifyShellLayer -def _unused_client_factory(): - raise AssertionError("shellctl client should not be used by these drive-layer tests") - - def _shell_layer() -> DifyShellLayer: return DifyShellLayer.from_config_with_settings( DifyShellLayerConfig(agent_stub_drive_ref="agent-1"), - shell_provider=ShellctlProvider(entrypoint="http://shellctl", token="", client_factory=_unused_client_factory), ) diff --git a/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py b/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py index b6ad74965b9..30405baff66 100644 --- a/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py +++ b/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py @@ -7,7 +7,6 @@ from dify_agent.layers.shell import ( DifyShellCliToolConfig, DifyShellEnvVarConfig, DifyShellLayerConfig, - DifyShellSandboxConfig, DifyShellSecretRefConfig, ) @@ -18,7 +17,6 @@ def test_shell_package_exports_client_safe_config_symbols_only() -> None: "DifyShellCliToolConfig", "DifyShellEnvVarConfig", "DifyShellLayerConfig", - "DifyShellSandboxConfig", "DifyShellSecretRefConfig", ] assert DIFY_SHELL_LAYER_TYPE_ID == "dify.shell" @@ -33,7 +31,6 @@ def test_shell_layer_config_defaults_and_forbids_unknown_fields() -> None: "cli_tools": [], "env": [], "secret_refs": [], - "sandbox": None, "redact_patterns": [], } @@ -54,7 +51,6 @@ def test_shell_layer_config_accepts_agent_soul_shell_settings() -> None: env=[DifyShellEnvVarConfig(name="PROJECT_NAME", value="demo")], secret_refs=[DifyShellSecretRefConfig(name="OPENAI_API_KEY", ref="credential-1")], agent_stub_drive_ref="agent-1", - sandbox=DifyShellSandboxConfig(provider="independent", config={"cpu": 2}), ) assert config.cli_tools[0].install_commands == ["apt-get update", "apt-get install -y ripgrep"] @@ -63,8 +59,6 @@ def test_shell_layer_config_accepts_agent_soul_shell_settings() -> None: assert config.env[0].name == "PROJECT_NAME" assert config.secret_refs[0].ref == "credential-1" assert config.agent_stub_drive_ref == "agent-1" - assert config.sandbox is not None - assert config.sandbox.config == {"cpu": 2} def test_shell_layer_config_rejects_invalid_env_names() -> None: diff --git a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py index 3d1e8f292cd..a358d0725ce 100644 --- a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py @@ -4,17 +4,22 @@ import asyncio import json from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from pathlib import Path from typing import cast import pytest import dify_agent.layers.shell.layer as shell_layer_module -from dify_agent.layers.shell import DIFY_SHELL_LAYER_TYPE_ID, DifyShellEnvVarConfig, DifyShellLayerConfig +from dify_agent.layers.shell import ( + DIFY_SHELL_LAYER_TYPE_ID, + DifyShellCliToolConfig, + DifyShellEnvVarConfig, + DifyShellLayerConfig, +) from dify_agent.layers.shell.layer import ( CompleteRemoteCommandResult, DEFAULT_TERMINATE_GRACE_SECONDS, DifyShellLayer, + DifyShellLayerDeps, DifyShellRuntimeState, ) from dify_agent.adapters.shell.protocols import ( @@ -22,11 +27,19 @@ from dify_agent.adapters.shell.protocols import ( ShellCommandStatus, ShellFileTransferProtocol, ShellProviderError, - SandboxExpiredError, - ShellProviderProtocol, - ShellResourceProtocol, ) from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig +from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer +from dify_agent.layers.runtime import DifyRuntimeLayerConfig +from dify_agent.layers.runtime.layer import DifyRuntimeLayer +from dify_agent.runtime_backend import ( + ExecutionBindingBackend, + RuntimeLayout, + RuntimeLease, + WorkspaceFileContent, + WorkspaceListResult, + WorkspaceReadResult, +) def _command_result( @@ -143,6 +156,16 @@ class _UnexpectedToolError(Exception): class FakeFiles(ShellFileTransferProtocol): + async def list_directory(self, *, path: str, limit: int) -> WorkspaceListResult: + raise AssertionError("resource.files should not be used by shell layer logic") + + async def read_file(self, *, path: str, max_bytes: int) -> WorkspaceReadResult: + raise AssertionError("resource.files should not be used by shell layer logic") + + async def read_bytes(self, *, path: str, max_bytes: int) -> WorkspaceFileContent: + del max_bytes + raise AssertionError("resource.files should not be used by shell layer logic") + async def upload(self, *, content: bytes, remote_path: str, cwd: str | None = None) -> None: raise AssertionError("resource.files should not be used by production shell layer logic") @@ -208,41 +231,23 @@ class FakeCommands: @dataclass(slots=True) -class FakeResource(ShellResourceProtocol): +class FakeResource: commands: FakeCommands files: FakeFiles = field(default_factory=FakeFiles) - suspended: bool = False - deleted: bool = False - _sandbox_id: str | None = None - - @property - def sandbox_id(self) -> str | None: - return self._sandbox_id - - async def suspend(self) -> None: - self.suspended = True - - async def delete(self) -> None: - self.deleted = True + handle: str = "sandbox-1" + layout: RuntimeLayout = field( + default_factory=lambda: RuntimeLayout( + home_dir="/home/agent-1", + workspace_dir="/home/agent-1/workspace/abc12ff", + ) + ) @dataclass(slots=True) -class FakeProvider(ShellProviderProtocol): +class FakeProvider: + """Test fixture retaining the old name while representing an active lease.""" + resource: FakeResource - create_calls: int = 0 - attach_calls: list[str] = field(default_factory=list) - attach_error: ShellProviderError | None = None - - async def create(self) -> ShellResourceProtocol: - self.create_calls += 1 - return self.resource - - async def attach(self, sandbox_id: str) -> ShellResourceProtocol: - self.attach_calls.append(sandbox_id) - if self.attach_error is not None: - raise self.attach_error - self.resource._sandbox_id = sandbox_id - return self.resource def _layer( @@ -252,22 +257,31 @@ def _layer( shell_home_root: str = "/home", ) -> tuple[DifyShellLayer, FakeProvider]: provider = FakeProvider(resource=FakeResource(commands=commands)) + root = shell_home_root.rstrip("/") + provider.resource.layout = RuntimeLayout( + home_dir=f"{root}/agent-1", + workspace_dir=f"{root}/agent-1/workspace/abc12ff", + ) layer = DifyShellLayer.from_config_with_settings( config or DifyShellLayerConfig(), - shell_provider=provider, - shell_home_root=shell_home_root, + ) + runtime = DifyRuntimeLayer.from_config_with_backend( + DifyRuntimeLayerConfig(backend_binding_ref="binding-1"), + backend=cast(ExecutionBindingBackend, object()), + ) + runtime._lease = cast(RuntimeLease, cast(object, provider.resource)) + layer.deps = DifyShellLayerDeps( + runtime=runtime, + execution_context=None, ) return layer, provider -@dataclass(slots=True) -class _ExecutionContextStub: - config: DifyExecutionContextLayerConfig - - def _bind_execution_context(layer: DifyShellLayer, *, agent_id: str | None = "agent-1") -> None: - layer.deps.execution_context = cast( - object, _ExecutionContextStub(config=_execution_context_config(agent_id=agent_id)) + layer.deps.execution_context = DifyExecutionContextLayer.from_config_with_settings( + _execution_context_config(agent_id=agent_id), + daemon_url="http://plugin-daemon", + daemon_api_key="", ) @@ -290,10 +304,8 @@ def _runtime_state( job_ids: list[str] | None = None, job_offsets: dict[str, int] | None = None, ) -> DifyShellRuntimeState: + del session_id, workspace_cwd, sandbox_id return DifyShellRuntimeState( - session_id=session_id, - workspace_cwd=workspace_cwd, - sandbox_id=sandbox_id, job_ids=[] if job_ids is None else job_ids, job_offsets={} if job_offsets is None else job_offsets, ) @@ -303,31 +315,12 @@ def test_shell_type_id_constant_matches_implementation_class() -> None: assert DIFY_SHELL_LAYER_TYPE_ID == DifyShellLayer.type_id -def test_resource_context_calls_provider_create_and_suspends_on_exit() -> None: - layer, provider = _layer(commands=FakeCommands()) - - async def scenario() -> None: - async with layer.resource_context(): - assert provider.create_calls == 1 - assert provider.resource.suspended is False - assert provider.resource.deleted is False - assert provider.resource.suspended is True - assert provider.resource.deleted is False - - asyncio.run(scenario()) - - -def test_shell_layer_create_allocates_workspace_and_bootstraps(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(shell_layer_module.time, "time", lambda: int("abc12", 16)) - monkeypatch.setattr(shell_layer_module.secrets, "token_hex", lambda _nbytes: "ff") +def test_shell_layer_create_bootstraps_inside_sandbox_workspace() -> None: expected_home = "/home/agent-1" expected_workspace_cwd = "/home/agent-1/workspace/abc12ff" def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: assert env == {"HOME": expected_home} - if cwd is None: - assert 'mkdir -p "$HOME/workspace"' in script - return _command_result("mkdir-job", status="exited", done=True, exit_code=0) assert cwd == expected_workspace_cwd assert "apt-get install -y ripgrep" in script return _command_result("bootstrap-job", status="exited", done=True, exit_code=0) @@ -335,7 +328,7 @@ def test_shell_layer_create_allocates_workspace_and_bootstraps(monkeypatch: pyte layer, provider = _layer( commands=FakeCommands(run_handler=run_handler), config=DifyShellLayerConfig( - cli_tools=[{"name": "ripgrep", "install_commands": ["apt-get install -y ripgrep"]}], + cli_tools=[DifyShellCliToolConfig(name="ripgrep", install_commands=["apt-get install -y ripgrep"])], ), ) _bind_execution_context(layer) @@ -343,31 +336,18 @@ def test_shell_layer_create_allocates_workspace_and_bootstraps(monkeypatch: pyte async def scenario() -> None: async with layer.resource_context(): await layer.on_context_create() - assert provider.resource.suspended is False asyncio.run(scenario()) - assert layer.runtime_state.session_id == "abc12ff" - assert layer.runtime_state.workspace_cwd == "~/workspace/abc12ff" - assert [call.job_id for call in provider.resource.commands.delete_calls] == ["mkdir-job", "bootstrap-job"] - assert provider.resource.suspended is True + assert [call.job_id for call in provider.resource.commands.delete_calls] == ["bootstrap-job"] -def test_shell_layer_uses_agent_specific_home_and_workspace_cwd( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(shell_layer_module.time, "time", lambda: int("abc12", 16)) - monkeypatch.setattr(shell_layer_module.secrets, "token_hex", lambda _nbytes: "ff") - +def test_shell_layer_uses_sandbox_layout_for_home_and_workspace_cwd() -> None: expected_home = "/home/agent-1" expected_workspace_cwd = "/home/agent-1/workspace/abc12ff" def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: del timeout - if script.startswith('mkdir -p "$HOME/workspace";'): - assert cwd is None - assert env == {"HOME": expected_home} - return _command_result("mkdir-job", status="exited", done=True, exit_code=0) if script == "pwd": assert cwd == expected_workspace_cwd assert env == {"HOME": expected_home} @@ -381,7 +361,6 @@ def test_shell_layer_uses_agent_specific_home_and_workspace_cwd( async def scenario() -> None: async with layer.resource_context(): - await layer.on_context_create() run_result = await tools["shell_run"].function_schema.call({"script": "pwd"}, None) # pyright: ignore[reportArgumentType] metadata, output = _parse_tagged_observation(run_result) assert metadata["job_id"] == "user-job" @@ -389,75 +368,26 @@ def test_shell_layer_uses_agent_specific_home_and_workspace_cwd( asyncio.run(scenario()) - assert layer.runtime_state.session_id == "abc12ff" - assert layer.runtime_state.workspace_cwd == "~/workspace/abc12ff" assert layer.runtime_state.job_ids == ["user-job"] - assert [call.job_id for call in commands.delete_calls] == ["mkdir-job"] + assert commands.delete_calls == [] -def test_shell_layer_uses_configured_home_root_for_local_development( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setattr(shell_layer_module.time, "time", lambda: int("abc12", 16)) - monkeypatch.setattr(shell_layer_module.secrets, "token_hex", lambda _nbytes: "ff") - - shell_home_root = tmp_path / "shell-home" - expected_home = f"{shell_home_root}/agent-1" - expected_workspace_cwd = f"{expected_home}/workspace/abc12ff" - - def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: - del timeout - if script.startswith('mkdir -p "$HOME/workspace";'): - assert cwd is None - assert env == {"HOME": expected_home} - return _command_result("mkdir-job", status="exited", done=True, exit_code=0) - if script == "pwd": - assert cwd == expected_workspace_cwd - assert env == {"HOME": expected_home} - return _command_result("user-job", status="exited", done=True, exit_code=0, output=expected_home, offset=13) - raise AssertionError(f"Unexpected script: {script!r}") - - layer, _provider = _layer(commands=FakeCommands(run_handler=run_handler), shell_home_root=f"{shell_home_root}/") - _bind_execution_context(layer) - tools = {tool.name: tool for tool in layer.tools} - - async def scenario() -> None: - async with layer.resource_context(): - await layer.on_context_create() - await tools["shell_run"].function_schema.call({"script": "pwd"}, None) # pyright: ignore[reportArgumentType] - - asyncio.run(scenario()) - - assert layer.shell_home_root == str(shell_home_root) - assert layer.runtime_state.workspace_cwd == "~/workspace/abc12ff" - - -def test_shell_layer_suspend_does_not_close_before_resource_context_exits() -> None: - layer, provider = _layer(commands=FakeCommands()) +def test_shell_layer_suspend_cleans_tracked_jobs_without_owning_sandbox() -> None: + commands = FakeCommands() + layer, _provider = _layer(commands=commands) layer.runtime_state = _runtime_state() async def scenario() -> None: async with layer.resource_context(): await layer.on_context_suspend() - assert provider.resource.suspended is False - assert provider.resource.suspended is True asyncio.run(scenario()) + assert commands.delete_calls == [] -def test_shell_layer_resume_recreates_live_home_and_workspace() -> None: - expected_home = "/home/agent-1" - - def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: - del timeout - assert script == 'mkdir -p "$HOME/workspace/abc12ff"' - assert cwd is None - assert env == {"HOME": expected_home} - return _command_result("resume-job", status="exited", done=True, exit_code=0) - - commands = FakeCommands(run_handler=run_handler) - layer, provider = _layer(commands=commands) +def test_shell_layer_resume_requires_active_sandbox_lease_only() -> None: + commands = FakeCommands() + layer, _provider = _layer(commands=commands) _bind_execution_context(layer) layer.runtime_state = _runtime_state() @@ -467,34 +397,12 @@ def test_shell_layer_resume_recreates_live_home_and_workspace() -> None: asyncio.run(scenario()) - assert [call.job_id for call in commands.delete_calls] == ["resume-job"] - assert provider.resource.suspended is True - - -def test_shell_layer_resume_requires_agent_id_for_live_workspace_reentry() -> None: - commands = FakeCommands() - layer, _provider = _layer(commands=commands) - _bind_execution_context(layer, agent_id=None) - layer.runtime_state = _runtime_state() - - async def scenario() -> None: - async with layer.resource_context(): - with pytest.raises(ValueError, match="requires execution_context\\.agent_id"): - await layer.on_context_resume() - - asyncio.run(scenario()) assert commands.run_calls == [] -def test_shell_layer_delete_cleans_workspace_and_tracked_jobs() -> None: - def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: - del env, timeout - assert cwd is None - assert script == 'rm -rf -- "$HOME/workspace/abc12ff"' - return _command_result("cleanup-job", status="exited", done=True, exit_code=0) - - commands = FakeCommands(run_handler=run_handler) - layer, provider = _layer(commands=commands) +def test_shell_layer_delete_cleans_tracked_jobs_without_deleting_workspace() -> None: + commands = FakeCommands() + layer, _provider = _layer(commands=commands) _bind_execution_context(layer) layer.runtime_state = _runtime_state(job_ids=["user-job"], job_offsets={"user-job": 9}) @@ -504,10 +412,9 @@ def test_shell_layer_delete_cleans_workspace_and_tracked_jobs() -> None: asyncio.run(scenario()) - assert [call.job_id for call in commands.delete_calls] == ["cleanup-job", "user-job"] + assert [call.job_id for call in commands.delete_calls] == ["user-job"] assert layer.runtime_state.job_ids == [] assert layer.runtime_state.job_offsets == {} - assert provider.resource.deleted is True def test_shell_layer_tools_map_inputs_and_maintain_offsets_with_tail_end() -> None: @@ -559,7 +466,7 @@ def test_shell_layer_tools_map_inputs_and_maintain_offsets_with_tail_end() -> No tail_handler=tail_handler, interrupt_handler=interrupt_handler, ) - layer, provider = _layer(commands=commands) + layer, _provider = _layer(commands=commands) _bind_execution_context(layer) tools = {tool.name: tool for tool in layer.tools} layer.runtime_state = _runtime_state() @@ -611,7 +518,6 @@ def test_shell_layer_tools_map_inputs_and_maintain_offsets_with_tail_end() -> No assert layer.runtime_state.job_offsets == {"user-job": 34} assert commands.tail_calls == [TailCall(job_id="user-job"), TailCall(job_id="user-job")] - assert provider.resource.suspended is True def test_shell_run_keeps_original_offset_when_tail_lookup_fails_for_truncated_output() -> None: @@ -633,7 +539,7 @@ def test_shell_run_keeps_original_offset_when_tail_lookup_fails_for_truncated_ou raise RuntimeError(f"tail unavailable for {job_id}") commands = FakeCommands(run_handler=run_handler, tail_handler=tail_handler) - layer, provider = _layer(commands=commands) + layer, _provider = _layer(commands=commands) _bind_execution_context(layer) tools = {tool.name: tool for tool in layer.tools} layer.runtime_state = _runtime_state() @@ -656,7 +562,6 @@ def test_shell_run_keeps_original_offset_when_tail_lookup_fails_for_truncated_ou assert layer.runtime_state.job_offsets == {"user-job": 10} assert commands.tail_calls == [TailCall(job_id="user-job")] - assert provider.resource.suspended is True def test_shell_run_formats_large_non_truncated_output_without_tail_lookup() -> None: @@ -831,7 +736,10 @@ def test_shell_run_returns_error_observation_and_logs_unexpected_exception( asyncio.run(scenario()) assert logged == [ - ("Unexpected shell tool failure: tool=%s session_id=%s job_id=%s", ("shell_run", "abc12ff", None)) + ( + "Unexpected shell tool failure: tool=%s job_id=%s", + ("shell_run", None), + ) ] @@ -856,7 +764,10 @@ def test_shell_wait_returns_error_observation_and_logs_unexpected_exception( asyncio.run(scenario()) assert logged == [ - ("Unexpected shell tool failure: tool=%s session_id=%s job_id=%s", ("shell_wait", "abc12ff", "user-job")) + ( + "Unexpected shell tool failure: tool=%s job_id=%s", + ("shell_wait", "user-job"), + ) ] @@ -884,8 +795,8 @@ def test_shell_input_returns_error_observation_and_logs_unexpected_exception( asyncio.run(scenario()) assert logged == [ ( - "Unexpected shell tool failure: tool=%s session_id=%s job_id=%s", - ("shell_input", "abc12ff", "user-job"), + "Unexpected shell tool failure: tool=%s job_id=%s", + ("shell_input", "user-job"), ) ] @@ -914,8 +825,8 @@ def test_shell_interrupt_returns_error_observation_and_logs_unexpected_exception asyncio.run(scenario()) assert logged == [ ( - "Unexpected shell tool failure: tool=%s session_id=%s job_id=%s", - ("shell_interrupt", "abc12ff", "user-job"), + "Unexpected shell tool failure: tool=%s job_id=%s", + ("shell_interrupt", "user-job"), ) ] @@ -947,12 +858,48 @@ def test_shell_interrupt_logs_unexpected_tail_failure_but_still_succeeds( asyncio.run(scenario()) assert logged == [ ( - "Failed to fetch output path for interrupted shell job %s in session %s", - ("user-job", "abc12ff"), + "Failed to fetch output path for interrupted shell job %s", + ("user-job",), ) ] +def test_agent_stub_token_omits_workspace_session_identity() -> None: + seen_session_ids: list[str | None] = [] + + def token_factory( + execution_context: DifyExecutionContextLayerConfig, + *, + session_id: str | None, + ) -> str: + del execution_context + seen_session_ids.append(session_id) + return "stub-token" + + def run_handler( + _script: str, + _cwd: str | None, + env: Mapping[str, str] | None, + _timeout: float, + ) -> ShellCommandResult: + assert env is not None + assert env["DIFY_AGENT_STUB_AUTH_JWE"] == "stub-token" + return _command_result("remote-job", status="exited", done=True, exit_code=0) + + layer, _provider = _layer(commands=FakeCommands(run_handler=run_handler)) + _bind_execution_context(layer) + layer.agent_stub_api_base_url = "http://localhost:5050/agent-stub" + layer.agent_stub_token_factory = token_factory + + async def scenario() -> None: + async with layer.resource_context(): + _ = await layer.run_remote_script_complete("true", inject_agent_stub_env=True) + + asyncio.run(scenario()) + + assert seen_session_ids == [None] + + def test_shell_run_propagates_cancelled_error() -> None: commands = FakeCommands( run_handler=lambda script, cwd, env, timeout: (_ for _ in ()).throw(asyncio.CancelledError()) @@ -1217,59 +1164,22 @@ def test_shell_layer_rejects_untracked_job_ids_without_provider_calls() -> None: assert commands.interrupt_calls == [] -def test_shell_layer_requires_agent_id_for_live_command_execution(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(shell_layer_module.time, "time", lambda: int("abc12", 16)) - monkeypatch.setattr(shell_layer_module.secrets, "token_hex", lambda _nbytes: "ff") - commands = FakeCommands() - layer, _provider = _layer(commands=commands) - _bind_execution_context(layer, agent_id=None) - - async def scenario() -> None: - async with layer.resource_context(): - with pytest.raises(ValueError, match="requires execution_context\\.agent_id"): - await layer.on_context_create() - - asyncio.run(scenario()) - assert commands.run_calls == [] - - def test_shell_layer_hooks_and_tools_fail_clearly_outside_active_resource_context() -> None: layer, _provider = _layer(commands=FakeCommands()) + layer.deps.runtime._lease = None tools = {tool.name: tool for tool in layer.tools} layer.runtime_state = _runtime_state() async def scenario() -> None: result = await tools["shell_run"].function_schema.call({"script": "pwd"}, None) # pyright: ignore[reportArgumentType] - _assert_error_observation(result, includes="shell resource") + _assert_error_observation(result, includes="resource_context") asyncio.run(scenario()) -def test_shell_runtime_state_validates_workspace_identity_and_offset_keys() -> None: - with pytest.raises(ValueError, match="5\\+2 lowercase hex format"): - _ = DifyShellRuntimeState.model_validate( - { - "session_id": "../../tmp", - "workspace_cwd": "~/workspace/../../tmp", - "job_ids": [], - "job_offsets": {}, - } - ) - - with pytest.raises(ValueError, match="workspace_cwd must equal"): - _ = DifyShellRuntimeState.model_validate( - { - "session_id": "abc12ff", - "workspace_cwd": "~/workspace/def34aa", - "job_ids": [], - "job_offsets": {}, - } - ) - +def test_shell_runtime_state_validates_offset_keys() -> None: state = DifyShellRuntimeState.model_validate( { - "session_id": "abc12ff", - "workspace_cwd": "~/workspace/abc12ff", "job_ids": ['job"bad with spaces'], "job_offsets": {'job"bad with spaces': 0}, } @@ -1278,108 +1188,12 @@ def test_shell_runtime_state_validates_workspace_identity_and_offset_keys() -> N with pytest.raises(ValueError, match="unknown job ids"): _ = DifyShellRuntimeState.model_validate( { - "session_id": "abc12ff", - "workspace_cwd": "~/workspace/abc12ff", "job_ids": ["job-1"], "job_offsets": {"job-2": 3}, } ) -def test_resource_context_attaches_when_sandbox_id_is_present() -> None: - layer, provider = _layer(commands=FakeCommands()) - layer.runtime_state = _runtime_state(sandbox_id="existing-sandbox-1") - - async def scenario() -> None: - async with layer.resource_context(): - assert provider.create_calls == 0 - assert provider.attach_calls == ["existing-sandbox-1"] - assert provider.resource.suspended is False - assert provider.resource.suspended is True - - asyncio.run(scenario()) - - -def test_resource_context_deletes_on_context_delete() -> None: - commands = FakeCommands( - run_handler=lambda script, cwd, env, timeout: _command_result( - "cleanup-job", status="exited", done=True, exit_code=0 - ) - ) - layer, provider = _layer(commands=commands) - _bind_execution_context(layer) - layer.runtime_state = _runtime_state() - - async def scenario() -> None: - async with layer.resource_context(): - await layer.on_context_delete() - assert provider.resource.deleted is True - assert provider.resource.suspended is False - - asyncio.run(scenario()) - - -def test_resource_context_suspends_on_context_suspend() -> None: - layer, provider = _layer(commands=FakeCommands()) - layer.runtime_state = _runtime_state() - - async def scenario() -> None: - async with layer.resource_context(): - await layer.on_context_suspend() - assert provider.resource.suspended is True - assert provider.resource.deleted is False - - asyncio.run(scenario()) - - -def test_resource_context_persists_sandbox_id_from_provider() -> None: - layer, provider = _layer(commands=FakeCommands()) - provider.resource._sandbox_id = "new-sandbox-42" - - async def scenario() -> None: - async with layer.resource_context(): - pass - - asyncio.run(scenario()) - assert layer.runtime_state.sandbox_id == "new-sandbox-42" - - -def test_resource_context_propagates_sandbox_expired() -> None: - """When attach() raises SandboxExpiredError, the error propagates to the caller. - The user must start a new session — no in-place recovery is attempted.""" - layer, provider = _layer(commands=FakeCommands()) - provider.attach_error = SandboxExpiredError( - "stale-sandbox-1", - cause=ShellProviderError( - 'request_failed (404): {"reason":"NOT_FOUND","message":"error: code = 400 reason = sandbox_expired message = sandbox has expired"}', - code="request_failed", - ), - ) - layer.runtime_state = _runtime_state(sandbox_id="stale-sandbox-1") - - async def scenario() -> None: - with pytest.raises(SandboxExpiredError): - async with layer.resource_context(): - pass - - asyncio.run(scenario()) - assert provider.attach_calls == ["stale-sandbox-1"] - assert provider.create_calls == 0 - - -def test_resource_context_reraises_non_expired_attach_error() -> None: - layer, provider = _layer(commands=FakeCommands()) - provider.attach_error = ShellProviderError("some other error", code="request_failed") - layer.runtime_state = _runtime_state(sandbox_id="sandbox-1") - - async def scenario() -> None: - async with layer.resource_context(): - pass - - with pytest.raises(ShellProviderError, match="some other error"): - asyncio.run(scenario()) - - # --------------------------------------------------------------------------- # Output redaction tests # --------------------------------------------------------------------------- @@ -1393,15 +1207,10 @@ def _layer_with_redaction( token_value: str = "eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0.fake-long-jwe-token-value", ) -> tuple[DifyShellLayer, FakeProvider]: """Create a layer with agent_stub env injection and optional redaction patterns.""" - provider = FakeProvider(resource=FakeResource(commands=commands)) - layer = DifyShellLayer.from_config_with_settings( - config or DifyShellLayerConfig(), - shell_provider=provider, - shell_home_root="/home", - shell_redact_patterns=shell_redact_patterns, - agent_stub_api_base_url="http://localhost:5050/agent-stub", - agent_stub_token_factory=lambda execution_context, session_id: token_value, - ) + layer, provider = _layer(commands=commands, config=config) + layer.shell_redact_patterns = shell_redact_patterns or [] + layer.agent_stub_api_base_url = "http://localhost:5050/agent-stub" + layer.agent_stub_token_factory = lambda execution_context, session_id: token_value return layer, provider diff --git a/dify-agent/tests/local/dify_agent/layers/test_runtime_layer.py b/dify-agent/tests/local/dify_agent/layers/test_runtime_layer.py new file mode 100644 index 00000000000..cdab88e90f2 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/layers/test_runtime_layer.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import cast + +import pytest + +from dify_agent.layers.runtime import DifyRuntimeLayer, DifyRuntimeLayerConfig +from dify_agent.runtime_backend import RuntimeLayout, RuntimeLease + + +@dataclass(slots=True) +class _Backend: + lease: RuntimeLease + acquired: list[str] = field(default_factory=list) + released: list[RuntimeLease] = field(default_factory=list) + + async def acquire(self, binding_ref: str) -> RuntimeLease: + self.acquired.append(binding_ref) + return self.lease + + async def release(self, lease: RuntimeLease) -> None: + self.released.append(lease) + + +@pytest.mark.anyio +async def test_runtime_layer_acquires_and_releases_operation_scoped_lease() -> None: + lease = cast( + RuntimeLease, + type( + "Lease", + (), + { + "layout": RuntimeLayout(home_dir="/home/agent", workspace_dir="/workspace"), + "commands": object(), + "files": object(), + }, + )(), + ) + backend = _Backend(lease=lease) + layer = DifyRuntimeLayer.from_config_with_backend( + DifyRuntimeLayerConfig(backend_binding_ref="binding-1"), + backend=backend, # pyright: ignore[reportArgumentType] + ) + + async with layer.resource_context(): + assert layer.lease is lease + await layer.on_context_create() + await layer.on_context_suspend() + await layer.on_context_delete() + + assert backend.acquired == ["binding-1"] + assert backend.released == [lease] + with pytest.raises(RuntimeError, match="resource_context"): + _ = layer.lease + + +def test_runtime_layer_config_contains_only_backend_binding_ref() -> None: + config = DifyRuntimeLayerConfig(backend_binding_ref="binding-1") + + assert config.model_dump() == {"backend_binding_ref": "binding-1"} diff --git a/dify-agent/tests/local/dify_agent/protocol/test_sandbox_locator.py b/dify-agent/tests/local/dify_agent/protocol/test_sandbox_locator.py deleted file mode 100644 index 6de05187708..00000000000 --- a/dify-agent/tests/local/dify_agent/protocol/test_sandbox_locator.py +++ /dev/null @@ -1,178 +0,0 @@ -from __future__ import annotations - -import pytest -from agenton.compositor import CompositorSessionSnapshot -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from agenton_collections.layers.plain import PromptLayerConfig -from dify_agent.layers.dify_core_tools import DIFY_CORE_TOOLS_LAYER_TYPE_ID -from dify_agent.layers.dify_plugin import DIFY_PLUGIN_LLM_LAYER_TYPE_ID -from dify_agent.layers.execution_context import DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, DifyExecutionContextLayerConfig -from dify_agent.layers.shell import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig -from dify_agent.protocol import ( - CreateRunRequest, - RunComposition, - RunLayerSpec, - RuntimeLayerSpec, - build_sandbox_locator_from_layer_specs, - build_sandbox_locator_from_run_request, - extract_runtime_layer_specs, -) - - -def _request() -> CreateRunRequest: - composition = RunComposition( - layers=[ - RunLayerSpec(name="prompt", type="plain.prompt", config=PromptLayerConfig(prefix="hi")), - RunLayerSpec( - name="execution_context", - type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, - config=DifyExecutionContextLayerConfig( - tenant_id="tenant-1", - user_from="account", - agent_mode="workflow_run", - invoke_from="service-api", - ), - ), - RunLayerSpec(name="llm", type=DIFY_PLUGIN_LLM_LAYER_TYPE_ID), - RunLayerSpec( - name="shell", - type=DIFY_SHELL_LAYER_TYPE_ID, - deps={"execution_context": "execution_context"}, - config=DifyShellLayerConfig(), - ), - ] - ) - snapshot = CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot(name="prompt", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), - LayerSessionSnapshot( - name="execution_context", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - LayerSessionSnapshot(name="llm", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), - LayerSessionSnapshot( - name="shell", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={"session_id": "abc12ff", "workspace_cwd": "~/workspace/abc12ff"}, - ), - ] - ) - return CreateRunRequest(composition=composition, session_snapshot=snapshot) - - -def test_build_sandbox_locator_from_run_request_filters_to_execution_context_and_shell() -> None: - locator = build_sandbox_locator_from_run_request(_request()) - - assert [layer.name for layer in locator.composition.layers] == ["execution_context", "shell"] - assert [layer.type for layer in locator.composition.layers] == [ - DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, - DIFY_SHELL_LAYER_TYPE_ID, - ] - assert [layer.name for layer in locator.session_snapshot.layers] == ["execution_context", "shell"] - - -def test_build_sandbox_locator_from_run_request_rejects_missing_session_snapshot() -> None: - request = _request() - request.session_snapshot = None - - with pytest.raises(ValueError, match="session_snapshot"): - build_sandbox_locator_from_run_request(request) - - -def test_extract_runtime_layer_specs_drops_sensitive_plugin_layers() -> None: - specs = extract_runtime_layer_specs(_request().composition) - - assert [spec.name for spec in specs] == ["prompt", "execution_context", "shell"] - - -def test_build_sandbox_locator_from_layer_specs_rejects_missing_shell() -> None: - with pytest.raises(ValueError, match="shell"): - build_sandbox_locator_from_layer_specs( - layer_specs=[RuntimeLayerSpec(name="execution_context", type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID)], - session_snapshot=CompositorSessionSnapshot(layers=[]), - ) - - -def test_build_sandbox_locator_from_layer_specs_rejects_missing_snapshot_layer() -> None: - specs = [ - RuntimeLayerSpec(name="execution_context", type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID), - RuntimeLayerSpec(name="shell", type=DIFY_SHELL_LAYER_TYPE_ID, deps={"execution_context": "execution_context"}), - ] - - with pytest.raises(ValueError, match="session_snapshot"): - build_sandbox_locator_from_layer_specs( - layer_specs=specs, - session_snapshot=CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="execution_context", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={} - ) - ] - ), - ) - - -def test_build_sandbox_locator_from_layer_specs_rejects_shell_dep_mismatch() -> None: - specs = [ - RuntimeLayerSpec(name="execution_context", type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID), - RuntimeLayerSpec(name="shell", type=DIFY_SHELL_LAYER_TYPE_ID, deps={"execution_context": "wrong-layer"}), - ] - - with pytest.raises(ValueError, match="depend on the execution_context"): - build_sandbox_locator_from_layer_specs( - layer_specs=specs, - session_snapshot=CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="execution_context", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - LayerSessionSnapshot( - name="shell", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - ] - ), - ) - - -def test_build_sandbox_locator_from_layer_specs_rejects_order_mismatch() -> None: - specs = [ - RuntimeLayerSpec(name="execution_context", type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID), - RuntimeLayerSpec(name="shell", type=DIFY_SHELL_LAYER_TYPE_ID, deps={"execution_context": "execution_context"}), - ] - - with pytest.raises(ValueError, match="in order"): - build_sandbox_locator_from_layer_specs( - layer_specs=specs, - session_snapshot=CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot(name="shell", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), - LayerSessionSnapshot( - name="execution_context", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={}, - ), - ] - ), - ) - - -def test_build_sandbox_locator_from_layer_specs_rejects_sensitive_runtime_specs() -> None: - with pytest.raises(ValueError, match="sensitive"): - build_sandbox_locator_from_layer_specs( - layer_specs=[RuntimeLayerSpec(name="llm", type=DIFY_PLUGIN_LLM_LAYER_TYPE_ID)], - session_snapshot=CompositorSessionSnapshot(layers=[]), - ) - - -def test_build_sandbox_locator_from_layer_specs_rejects_sensitive_core_tool_runtime_specs() -> None: - with pytest.raises(ValueError, match="sensitive"): - build_sandbox_locator_from_layer_specs( - layer_specs=[RuntimeLayerSpec(name="core_tools", type=DIFY_CORE_TOOLS_LAYER_TYPE_ID)], - session_snapshot=CompositorSessionSnapshot(layers=[]), - ) diff --git a/dify-agent/tests/local/dify_agent/protocol/test_working_environment.py b/dify-agent/tests/local/dify_agent/protocol/test_working_environment.py new file mode 100644 index 00000000000..a41a80f20ed --- /dev/null +++ b/dify-agent/tests/local/dify_agent/protocol/test_working_environment.py @@ -0,0 +1,47 @@ +import pytest +from pydantic import ValidationError + +from dify_agent.protocol import ( + CreateExecutionBindingRequest, + CreateHomeSnapshotFromBindingRequest, + DestroyExecutionBindingRequest, + WorkspaceListRequest, +) + + +def test_execution_binding_request_uses_opaque_backend_refs() -> None: + request = CreateExecutionBindingRequest( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref="opaque-workspace", + home_snapshot_ref="opaque-home", + ) + + assert request.model_dump() == { + "tenant_id": "tenant-1", + "agent_id": "agent-1", + "binding_id": "binding-1", + "workspace_id": "workspace-1", + "existing_workspace_ref": "opaque-workspace", + "home_snapshot_ref": "opaque-home", + } + + +def test_destroy_workspace_requires_workspace_ref() -> None: + with pytest.raises(ValidationError, match="workspace_ref"): + DestroyExecutionBindingRequest(binding_ref="binding-1", destroy_workspace=True) + + +def test_snapshot_and_file_requests_locate_binding_directly() -> None: + snapshot = CreateHomeSnapshotFromBindingRequest( + tenant_id="tenant-1", + agent_id="agent-1", + home_snapshot_id="home-2", + backend_binding_ref="binding-ref", + ) + listing = WorkspaceListRequest(backend_binding_ref="binding-ref", path="~/files") + + assert snapshot.backend_binding_ref == "binding-ref" + assert listing.path == "~/files" diff --git a/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py b/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py index e0ce4e7a50f..11b40b20bc7 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py @@ -2,15 +2,6 @@ import sys import types from typing import cast -from pydantic import BaseModel - - -if "pydantic_settings" not in sys.modules: - pydantic_settings = types.ModuleType("pydantic_settings") - pydantic_settings.BaseSettings = BaseModel - pydantic_settings.SettingsConfigDict = dict - sys.modules["pydantic_settings"] = pydantic_settings - if "graphon.model_runtime.entities.llm_entities" not in sys.modules: graphon_module = types.ModuleType("graphon") model_runtime_module = types.ModuleType("graphon.model_runtime") @@ -84,13 +75,15 @@ if "jsonschema" not in sys.modules: sys.modules["jsonschema.protocols"] = jsonschema_protocols_module sys.modules["jsonschema.validators"] = jsonschema_validators_module -from dify_agent.adapters.shell.protocols import ShellProviderProtocol from dify_agent.layers.dify_core_tools import DIFY_CORE_TOOLS_LAYER_TYPE_ID, DifyCoreToolsLayerConfig from dify_agent.layers.dify_core_tools.layer import DifyCoreToolsLayer +from dify_agent.layers.runtime import DIFY_RUNTIME_LAYER_TYPE_ID, DifyRuntimeLayerConfig +from dify_agent.layers.runtime.layer import DifyRuntimeLayer from dify_agent.layers.shell import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.layers.shell.layer import DifyShellLayer from dify_agent.runtime.compositor_factory import create_default_layer_providers +from dify_agent.runtime_backend import ExecutionBindingBackend, HomeSnapshotBackend, RuntimeBackendProfile class FakeProvider: @@ -100,19 +93,28 @@ class FakeProvider: raise AssertionError("create should not be called by these tests") -def test_default_layer_providers_wire_provided_shell_provider() -> None: - fake_provider = FakeProvider() +def _runtime_backend_profile() -> RuntimeBackendProfile: + return RuntimeBackendProfile( + home_snapshots=cast(HomeSnapshotBackend, FakeProvider()), + execution_bindings=cast(ExecutionBindingBackend, FakeProvider()), + ) + + +def test_default_layer_providers_register_runtime_layer() -> None: + profile = _runtime_backend_profile() providers = create_default_layer_providers( - shell_provider=cast(ShellProviderProtocol, fake_provider), - shell_home_root="/tmp/dify-agent-home", + runtime_backend_profile=profile, ) shell_provider = next(provider for provider in providers if provider.type_id == DIFY_SHELL_LAYER_TYPE_ID) shell_layer = shell_provider.create_layer(DifyShellLayerConfig()) + runtime_provider = next(provider for provider in providers if provider.type_id == DIFY_RUNTIME_LAYER_TYPE_ID) + runtime_layer = runtime_provider.create_layer(DifyRuntimeLayerConfig(backend_binding_ref="binding-1")) assert isinstance(shell_layer, DifyShellLayer) - assert shell_layer.shell_provider is fake_provider - assert shell_layer.shell_home_root == "/tmp/dify-agent-home" + assert isinstance(runtime_layer, DifyRuntimeLayer) + assert runtime_layer.backend is profile.execution_bindings + assert {provider.type_id for provider in providers} >= {"dify.runtime", "dify.shell"} def test_default_layer_providers_forward_agent_stub_token_factory() -> None: @@ -127,7 +129,7 @@ def test_default_layer_providers_forward_agent_stub_token_factory() -> None: return f"token-for:{execution_context.tenant_id}:{session_id}" providers = create_default_layer_providers( - shell_provider=cast(ShellProviderProtocol, FakeProvider()), + runtime_backend_profile=_runtime_backend_profile(), agent_stub_api_base_url="https://agent.example.com/agent-stub", agent_stub_token_factory=build_agent_stub_token, ) diff --git a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py b/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py index ac72d0502f1..f33a1639683 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py @@ -8,8 +8,10 @@ import pytest from agenton.compositor import CompositorSessionSnapshot, LayerSessionSnapshot from agenton.layers import LifecycleState from agenton_collections.layers.plain import PromptLayerConfig +from dify_agent.layers.dify_plugin import DifyPluginLLMLayerConfig +from dify_agent.layers.execution_context import DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, DifyExecutionContextLayerConfig from dify_agent.layers.output import DIFY_OUTPUT_LAYER_TYPE_ID, DifyOutputLayerConfig -from dify_agent.protocol import DIFY_AGENT_OUTPUT_LAYER_ID +from dify_agent.protocol import DIFY_AGENT_MODEL_LAYER_ID, DIFY_AGENT_OUTPUT_LAYER_ID from dify_agent.protocol.schemas import ( CancelRunRequest, CreateRunRequest, @@ -27,7 +29,30 @@ def _request( *, output_config: Mapping[str, object] | DifyOutputLayerConfig | None = None, ) -> CreateRunRequest: - layers = [RunLayerSpec(name="prompt", type="plain.prompt", config=PromptLayerConfig(user=user))] + layers = [ + RunLayerSpec(name="prompt", type="plain.prompt", config=PromptLayerConfig(user=user)), + RunLayerSpec( + name="execution_context", + type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, + config=DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_from="account", + agent_mode="workflow_run", + invoke_from="service-api", + ), + ), + RunLayerSpec( + name=DIFY_AGENT_MODEL_LAYER_ID, + type="dify.plugin.llm", + deps={"execution_context": "execution_context"}, + config=DifyPluginLLMLayerConfig( + plugin_id="langgenius/openai", + model_provider="openai", + model="demo-model", + credentials={"api_key": "secret"}, + ), + ), + ] if output_config is not None: layers.append( RunLayerSpec( @@ -339,7 +364,17 @@ def test_create_run_accepts_closed_session_snapshot_and_runner_fails_asynchronou name="prompt", lifecycle_state=LifecycleState.CLOSED, runtime_state={}, - ) + ), + LayerSessionSnapshot( + name="execution_context", + lifecycle_state=LifecycleState.SUSPENDED, + runtime_state={}, + ), + LayerSessionSnapshot( + name=DIFY_AGENT_MODEL_LAYER_ID, + lifecycle_state=LifecycleState.SUSPENDED, + runtime_state={}, + ), ] ) diff --git a/dify-agent/tests/local/dify_agent/runtime/test_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_runner.py index dceab28ab98..fa45481bc34 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_runner.py @@ -30,9 +30,8 @@ from agenton_collections.layers.pydantic_ai import PYDANTIC_AI_HISTORY_LAYER_TYP from agenton_collections.layers.plain import PromptLayerConfig, ToolsLayer from dify_agent.layers.ask_human import DIFY_ASK_HUMAN_LAYER_TYPE_ID, DifyAskHumanLayerConfig from dify_agent.layers.execution_context import DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, DifyExecutionContextLayerConfig +from dify_agent.layers.runtime import DIFY_RUNTIME_LAYER_TYPE_ID, DifyRuntimeLayerConfig from dify_agent.layers.shell import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig -from dify_agent.adapters.shell.shellctl import ShellctlProvider -from dify_agent.layers.shell.layer import DifyShellLayer from dify_agent.layers.dify_plugin.configs import ( DIFY_PLUGIN_TOOLS_LAYER_TYPE_ID, DifyPluginLLMLayerConfig, @@ -72,6 +71,16 @@ from dify_agent.runtime.runner import ( RunSuccessOutcome, _run_failed_error_payload, ) +from dify_agent.runtime_backend import ( + ExecutionBindingAllocation, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotBackend, + RuntimeBackendProfile, + RuntimeLayout, + RuntimeLease, +) +from dify_agent.runtime_backend.shellctl import ShellctlRuntimeLease, create_shellctl_lease from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView @@ -135,6 +144,33 @@ class FakeRunnerShellctlClient: return DeleteJobResponse(job_id=job_id) +class FakeRunnerExecutionBindingBackend: + def __init__(self, client: FakeRunnerShellctlClient) -> None: + self.client = client + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + return ExecutionBindingAllocation(binding_ref=spec.binding_id, workspace_ref=spec.workspace_id) + + async def acquire(self, binding_ref: str) -> RuntimeLease: + return create_shellctl_lease( + handle=binding_ref, + layout=RuntimeLayout( + home_dir="/home/agent-1", + workspace_dir="/home/agent-1/workspace/abc12ff", + ), + entrypoint="http://shellctl", + token="", + client_factory=lambda: self.client, # pyright: ignore[reportArgumentType] + ) + + async def release(self, lease: RuntimeLease) -> None: + assert isinstance(lease, ShellctlRuntimeLease) + await lease.close() + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + del spec + + def test_run_failed_error_payload_preserves_plugin_rate_limit_error() -> None: exc = ModelHTTPError( 429, @@ -273,44 +309,6 @@ def _request( ) -def _lifecycle_only_request( - *, - on_exit: LayerExitSignals | None = None, - session_snapshot: CompositorSessionSnapshot | None = None, - deferred_tool_results: DeferredToolResultsPayload | None = None, -) -> CreateRunRequest: - snapshot = session_snapshot or CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot(name="prompt", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), - LayerSessionSnapshot(name="execution_context", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), - ] - ) - return CreateRunRequest( - composition=RunComposition( - layers=[ - RunLayerSpec( - name="prompt", - type="plain.prompt", - config=PromptLayerConfig(prefix="system", user="hello"), - ), - RunLayerSpec( - name="execution_context", - type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, - config=DifyExecutionContextLayerConfig( - tenant_id="tenant-1", - user_from="account", - agent_mode="workflow_run", - invoke_from="service-api", - ), - ), - ] - ), - session_snapshot=snapshot, - deferred_tool_results=deferred_tool_results, - on_exit=on_exit or LayerExitSignals(default=ExitIntent.DELETE), - ) - - def _recursive_output_schema() -> dict[str, object]: return { "type": "object", @@ -1655,20 +1653,11 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( monkeypatch.setattr(DifyPluginToolsLayer, "get_tools", fake_get_tools) monkeypatch.setattr("dify_agent.runtime.runner.create_agent", fake_create_agent) - shell_provider = LayerProvider.from_factory( - layer_type=DifyShellLayer, - create=lambda config: DifyShellLayer.from_config_with_settings( - DifyShellLayerConfig.model_validate(config), - shell_provider=ShellctlProvider( - entrypoint="http://shellctl", - token="", - client_factory=lambda: shell_client, - ), - ), + runtime_backend_profile = RuntimeBackendProfile( + home_snapshots=cast(HomeSnapshotBackend, object()), + execution_bindings=FakeRunnerExecutionBindingBackend(shell_client), ) - layer_providers = tuple( - provider for provider in create_default_layer_providers() if provider.type_id != DIFY_SHELL_LAYER_TYPE_ID - ) + (shell_provider,) + layer_providers = create_default_layer_providers(runtime_backend_profile=runtime_backend_profile) request = CreateRunRequest( composition=RunComposition( @@ -1684,15 +1673,21 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( config=DifyExecutionContextLayerConfig( tenant_id="tenant-1", agent_id="agent-1", + agent_config_version_id="config-1", user_from="account", agent_mode="workflow_run", invoke_from="service-api", ), ), + RunLayerSpec( + name="runtime", + type=DIFY_RUNTIME_LAYER_TYPE_ID, + config=DifyRuntimeLayerConfig(backend_binding_ref="binding-1"), + ), RunLayerSpec( name="shell", type=DIFY_SHELL_LAYER_TYPE_ID, - deps={"execution_context": "execution_context"}, + deps={"execution_context": "execution_context", "runtime": "runtime"}, config=DifyShellLayerConfig(), ), RunLayerSpec( @@ -1746,7 +1741,7 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( asyncio.run(scenario()) assert create_agent_called is False - assert shell_client.delete_calls == [("mkdir-job", True, None)] + assert shell_client.delete_calls == [] assert shell_client.closed is True assert [event.type for event in sink.events["run-shell-duplicate-tools"]] == ["run_started", "run_failed"] assert sink.statuses["run-shell-duplicate-tools"] == "failed" @@ -1961,110 +1956,25 @@ def test_runner_applies_on_exit_overrides_to_success_snapshot(monkeypatch: pytes } -def test_runner_lifecycle_only_cleanup_succeeds_without_model_and_emits_no_pydantic_ai_events() -> None: - request = _lifecycle_only_request() - sink = InMemoryRunEventSink() - - async def scenario() -> None: - async with httpx.AsyncClient() as client: - await AgentRunRunner( - sink=sink, - request=request, - run_id="run-lifecycle-only", - plugin_daemon_http_client=client, - dify_api_http_client=client, - ).run() - - asyncio.run(scenario()) - - events = sink.events["run-lifecycle-only"] - assert [event.type for event in events] == ["run_started", "run_succeeded"] - terminal = events[-1] - assert isinstance(terminal, RunSucceededEvent) - assert terminal.data.output is None - assert terminal.data.usage is None - assert {layer.name: layer.lifecycle_state for layer in terminal.data.session_snapshot.layers} == { - "prompt": LifecycleState.CLOSED, - "execution_context": LifecycleState.CLOSED, - } - - -def test_runner_lifecycle_only_requires_session_snapshot() -> None: +def test_runner_requires_model_layer() -> None: request = _request(llm_layer_name="not-llm") sink = InMemoryRunEventSink() async def scenario() -> None: async with httpx.AsyncClient() as client: - with pytest.raises(AgentRunValidationError, match="session_snapshot"): + with pytest.raises(AgentRunValidationError, match="Missing required"): await AgentRunRunner( sink=sink, request=request, - run_id="run-lifecycle-only-missing-snapshot", + run_id="run-missing-model", plugin_daemon_http_client=client, dify_api_http_client=client, ).run() asyncio.run(scenario()) - assert [event.type for event in sink.events["run-lifecycle-only-missing-snapshot"]] == ["run_started", "run_failed"] - assert sink.statuses["run-lifecycle-only-missing-snapshot"] == "failed" - - -def test_runner_lifecycle_only_rejects_deferred_tool_results() -> None: - request = _lifecycle_only_request( - deferred_tool_results=DeferredToolResultsPayload.model_validate({"calls": {"tool-call-1": {"ok": True}}}) - ) - sink = InMemoryRunEventSink() - - async def scenario() -> None: - async with httpx.AsyncClient() as client: - with pytest.raises(AgentRunValidationError, match="Deferred tool results"): - await AgentRunRunner( - sink=sink, - request=request, - run_id="run-lifecycle-only-deferred-results", - plugin_daemon_http_client=client, - dify_api_http_client=client, - ).run() - - asyncio.run(scenario()) - - assert [event.type for event in sink.events["run-lifecycle-only-deferred-results"]] == [ - "run_started", - "run_failed", - ] - assert sink.statuses["run-lifecycle-only-deferred-results"] == "failed" - - -def test_runner_lifecycle_only_exit_hook_failure_emits_run_failed_not_validation_error( - monkeypatch: pytest.MonkeyPatch, -) -> None: - request = _lifecycle_only_request() - sink = InMemoryRunEventSink() - - def _explode(_run: object, _signals: LayerExitSignals) -> None: - raise RuntimeError("delete hook failed") - - monkeypatch.setattr("dify_agent.runtime.runner.apply_layer_exit_signals", _explode) - - async def scenario() -> None: - async with httpx.AsyncClient() as client: - with pytest.raises(RuntimeError, match="delete hook failed"): - await AgentRunRunner( - sink=sink, - request=request, - run_id="run-lifecycle-only-exit-hook-failure", - plugin_daemon_http_client=client, - dify_api_http_client=client, - ).run() - - asyncio.run(scenario()) - - assert [event.type for event in sink.events["run-lifecycle-only-exit-hook-failure"]] == [ - "run_started", - "run_failed", - ] - assert sink.statuses["run-lifecycle-only-exit-hook-failure"] == "failed" + assert [event.type for event in sink.events["run-missing-model"]] == ["run_started", "run_failed"] + assert sink.statuses["run-missing-model"] == "failed" def test_runner_passes_output_layer_spec_to_agent_and_serializes_structured_result( @@ -2757,7 +2667,7 @@ def test_runner_rejects_closed_session_snapshot_as_validation_error() -> None: assert sink.statuses["run-closed-snapshot"] == "failed" -def test_runner_treats_missing_shell_entrypoint_as_validation_error() -> None: +def test_runner_treats_missing_runtime_dependency_as_validation_error() -> None: request = CreateRunRequest( composition=RunComposition( layers=[ @@ -2795,7 +2705,7 @@ def test_runner_treats_missing_shell_entrypoint_as_validation_error() -> None: async def scenario() -> None: async with httpx.AsyncClient() as client: - with pytest.raises(AgentRunValidationError, match="non-null shell provider"): + with pytest.raises(AgentRunValidationError, match="Dependency 'runtime' is required"): await AgentRunRunner( sink=sink, request=request, @@ -2880,9 +2790,7 @@ def test_runner_treats_invalid_shell_snapshot_offsets_as_validation_error() -> N run_id="run-invalid-shell-offset", plugin_daemon_http_client=client, dify_api_http_client=client, - layer_providers=create_default_layer_providers( - shell_provider=ShellctlProvider(entrypoint="http://shellctl", token=""), - ), + layer_providers=create_default_layer_providers(), ).run() asyncio.run(scenario()) diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py new file mode 100644 index 00000000000..ba33ad76a41 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import cast + +import pytest + +from dify_agent.runtime_backend import ( + BindingCreateError, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotCreateSpec, + InitializeHomeSnapshotSpec, + SharedWorkspaceUnsupportedError, + WorkspacePreservationUnsupportedError, +) +from dify_agent.runtime_backend.e2b import ( + E2BExecutionBindingBackend, + E2BHomeSnapshotBackend, + E2BRuntimeLease, +) +from dify_agent.runtime_backend.shellctl import ShellctlRuntimeLease + + +@dataclass(slots=True) +class _Files: + paths: set[str] = field(default_factory=set) + + async def make_dir(self, path: str) -> bool: + self.paths.add(path) + return True + + async def exists(self, path: str) -> bool: + return path in self.paths + + async def remove(self, path: str) -> None: + self.paths.discard(path) + + +@dataclass(slots=True) +class _Snapshot: + snapshot_id: str + names: list[str] = field(default_factory=list) + + +@dataclass(slots=True) +class _Sandbox: + sandbox_id: str + files: _Files = field(default_factory=_Files) + traffic_access_token: str | None = "traffic-token" + pauses: list[bool] = field(default_factory=list) + killed: int = 0 + snapshots: int = 0 + pause_error: Exception | None = None + + def get_host(self, port: int) -> str: + return f"{self.sandbox_id}-{port}.example.test" + + async def pause(self, keep_memory: bool = True) -> bool: + self.pauses.append(keep_memory) + if self.pause_error is not None: + raise self.pause_error + return True + + async def kill(self) -> bool: + self.killed += 1 + return True + + async def create_snapshot(self, name: str | None = None) -> _Snapshot: + del name + self.snapshots += 1 + return _Snapshot(snapshot_id=f"snapshot-{self.sandbox_id}-{self.snapshots}") + + +@dataclass(slots=True) +class _ControlPlane: + created: list[tuple[str, str]] = field(default_factory=list) + sandboxes: dict[str, _Sandbox] = field(default_factory=dict) + killed: list[str] = field(default_factory=list) + deleted_snapshots: list[str] = field(default_factory=list) + pause_error: Exception | None = None + + async def create(self, template: str, *, timeout: int, metadata: dict[str, str], on_timeout: str) -> _Sandbox: + del timeout + sandbox_id = f"sandbox-{len(self.sandboxes) + 1}" + sandbox = _Sandbox(sandbox_id=sandbox_id, pause_error=self.pause_error) + self.sandboxes[sandbox_id] = sandbox + self.created.append((template, on_timeout)) + assert metadata["dify.resource"] in {"home-snapshot-initialize", "runtime-sandbox"} + return sandbox + + async def connect(self, handle: str, *, timeout: int) -> _Sandbox: + del timeout + return self.sandboxes[handle] + + async def kill(self, handle: str) -> bool: + self.killed.append(handle) + return True + + async def delete_snapshot(self, snapshot_ref: str) -> bool: + self.deleted_snapshots.append(snapshot_ref) + return True + + +@pytest.mark.anyio +async def test_e2b_profile_uses_snapshot_as_runtime_template_and_couples_refs() -> None: + control = _ControlPlane() + snapshots = E2BHomeSnapshotBackend( + control_plane=control, # pyright: ignore[reportArgumentType] + template="prepared-template", + active_timeout_seconds=3600, + ) + bindings = E2BExecutionBindingBackend( + control_plane=control, # pyright: ignore[reportArgumentType] + active_timeout_seconds=3600, + ) + + snapshot_ref = await snapshots.initialize( + InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1") + ) + allocation = await bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref=snapshot_ref, + ) + ) + + assert control.created == [("prepared-template", "kill"), (snapshot_ref, "pause")] + assert allocation.binding_ref == allocation.workspace_ref + runtime = control.sandboxes[allocation.binding_ref] + assert runtime.files.paths == {"/home/dify/workspace"} + assert runtime.pauses == [True] + + await bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=allocation.binding_ref, + workspace_ref=allocation.workspace_ref, + destroy_workspace=True, + ) + ) + await snapshots.delete(snapshot_ref) + + assert control.killed == [allocation.binding_ref] + assert control.deleted_snapshots == [snapshot_ref] + + +@pytest.mark.anyio +async def test_e2b_rejects_shared_workspace_and_binding_only_destroy() -> None: + control = _ControlPlane() + backend = E2BExecutionBindingBackend( + control_plane=control, # pyright: ignore[reportArgumentType] + active_timeout_seconds=3600, + ) + spec = ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-2", + binding_id="binding-2", + workspace_id="workspace-1", + existing_workspace_ref="sandbox-1", + home_snapshot_ref="snapshot-1", + ) + + with pytest.raises(SharedWorkspaceUnsupportedError): + await backend.create_binding(spec) + assert control.created == [] + assert control.sandboxes == {} + with pytest.raises(WorkspacePreservationUnsupportedError): + await backend.destroy_binding(ExecutionBindingDestroySpec(binding_ref="sandbox-1", destroy_workspace=False)) + + +@pytest.mark.anyio +async def test_e2b_binding_create_kills_sandbox_when_initialization_fails() -> None: + control = _ControlPlane(pause_error=RuntimeError("pause failed")) + backend = E2BExecutionBindingBackend( + control_plane=control, # pyright: ignore[reportArgumentType] + active_timeout_seconds=3600, + ) + + with pytest.raises(BindingCreateError, match="pause failed"): + await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="snapshot-1", + ) + ) + + sandbox = next(iter(control.sandboxes.values())) + assert sandbox.killed == 1 + + +@pytest.mark.anyio +async def test_e2b_checkpoint_uses_exact_source_runtime() -> None: + control = _ControlPlane() + source_sandbox = _Sandbox(sandbox_id="source") + source = E2BRuntimeLease( + sandbox=source_sandbox, + data_plane=cast(ShellctlRuntimeLease, object()), + ) + backend = E2BHomeSnapshotBackend( + control_plane=control, # pyright: ignore[reportArgumentType] + template="prepared-template", + active_timeout_seconds=3600, + ) + + snapshot_ref = await backend.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), + source=source, + ) + + assert snapshot_ref == "snapshot-source-1" + assert source_sandbox.snapshots == 1 diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py new file mode 100644 index 00000000000..7ae011cd24a --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py @@ -0,0 +1,310 @@ +from __future__ import annotations + +import json +from collections.abc import Callable +from typing import cast + +import httpx2 as httpx +import pytest + +from dify_agent.runtime_backend import ( + BindingAcquireError, + BindingDestroyError, + BindingLostError, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotCreateSpec, + InitializeHomeSnapshotSpec, + RuntimeLease, + WorkspacePreservationUnsupportedError, +) +from dify_agent.runtime_backend.enterprise import ( + EnterpriseExecutionBindingBackend, + EnterpriseHomeSnapshotBackend, +) + + +def _mock_http( + monkeypatch: pytest.MonkeyPatch, + handler: Callable[[httpx.Request], httpx.Response], +) -> list[httpx.AsyncClient]: + original_async_client = httpx.AsyncClient + transport = httpx.MockTransport(handler) + clients: list[httpx.AsyncClient] = [] + + def create_transport(*, retries: int = 0) -> httpx.AsyncBaseTransport: + del retries + return transport + + monkeypatch.setattr(httpx, "AsyncHTTPTransport", create_transport) + + def create_client(*args: object, **kwargs: object) -> httpx.AsyncClient: + _ = kwargs.setdefault("transport", transport) + client = original_async_client(*args, **kwargs) + clients.append(client) + return client + + monkeypatch.setattr(httpx, "AsyncClient", create_client) + return clients + + +def _job_response(*, exit_code: int = 0) -> httpx.Response: + return httpx.Response( + 200, + json={ + "job_id": "job-1", + "done": True, + "status": "exited", + "exit_code": exit_code, + "output_path": "/tmp/output.log", + "output": "", + "offset": 0, + "truncated": False, + }, + ) + + +@pytest.mark.anyio +async def test_enterprise_acquire_exposes_canonical_layout_through_gateway_proxy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert request.headers["X-Sandbox-Id"] == "sandbox-1" + assert request.headers["X-Inner-Api-Key"] == "secret" + if request.method == "POST": + payload = cast(dict[str, object], json.loads(request.content)) + script = payload["script"] + assert isinstance(script, str) + assert "test -d /home/dify" in script + assert "test -d /home/dify/workspace" in script + return _job_response() + return httpx.Response(200, json={"job_id": "job-1"}) + + clients = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + proxy_timeout=90, + ) + + lease = await backend.acquire("sandbox-1") + + assert lease.layout.home_dir == "/home/dify" + assert lease.layout.workspace_dir == "/home/dify/workspace" + assert [request.url.path for request in requests] == [ + "/proxy/v1/jobs/run", + "/proxy/v1/jobs/job-1", + ] + assert clients[0].timeout.read == 90 + await backend.release(lease) + assert clients[0].is_closed + + +@pytest.mark.anyio +@pytest.mark.parametrize( + ("status_code", "code", "expected_error"), + [ + (404, "sandbox_expired", BindingLostError), + (502, "upstream_failure", BindingAcquireError), + ], +) +async def test_enterprise_acquire_maps_proxy_failures_and_closes_transport( + monkeypatch: pytest.MonkeyPatch, + status_code: int, + code: str, + expected_error: type[Exception], +) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(status_code, json={"error": {"code": code, "message": "unavailable"}}) + + clients = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + ) + + with pytest.raises(expected_error): + _ = await backend.acquire("sandbox-1") + + assert clients[0].is_closed + + +@pytest.mark.anyio +async def test_enterprise_acquire_treats_missing_runtime_directories_as_lost( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return _job_response(exit_code=1) + return httpx.Response(200, json={"job_id": "job-1"}) + + clients = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + ) + + with pytest.raises(BindingLostError, match="Home or Workspace"): + _ = await backend.acquire("sandbox-1") + + assert clients[0].is_closed + + +@pytest.mark.anyio +async def test_enterprise_release_rejects_foreign_runtime_lease() -> None: + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + ) + + with pytest.raises(TypeError, match="only release its own RuntimeLease"): + await backend.release(cast(RuntimeLease, object())) + + +@pytest.mark.anyio +async def test_enterprise_destroy_validates_coupled_workspace_before_gateway_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(204) + + _ = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + ) + + with pytest.raises(WorkspacePreservationUnsupportedError): + await backend.destroy_binding(ExecutionBindingDestroySpec(binding_ref="sandbox-1", destroy_workspace=False)) + with pytest.raises(BindingDestroyError, match="must equal"): + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="sandbox-1", + workspace_ref="workspace-1", + destroy_workspace=True, + ) + ) + + assert requests == [] + + +@pytest.mark.anyio +@pytest.mark.parametrize("status_code", [204, 404]) +async def test_enterprise_destroy_is_authenticated_and_idempotent( + monkeypatch: pytest.MonkeyPatch, + status_code: int, +) -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(status_code) + + clients = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + ) + + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="sandbox-1", + workspace_ref="sandbox-1", + destroy_workspace=True, + ) + ) + + assert len(requests) == 1 + assert requests[0].method == "DELETE" + assert requests[0].url.path == "/v1/sandboxes/sandbox-1" + assert requests[0].headers["X-Inner-Api-Key"] == "secret" + assert clients[0].is_closed + + +@pytest.mark.anyio +async def test_enterprise_destroy_encodes_binding_ref_as_one_path_segment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(204) + + _ = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + ) + opaque_ref = "../admin?x=1" + + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=opaque_ref, + workspace_ref=opaque_ref, + destroy_workspace=True, + ) + ) + + assert len(requests) == 1 + assert requests[0].url.raw_path == b"/v1/sandboxes/..%2Fadmin%3Fx%3D1" + + +@pytest.mark.anyio +async def test_enterprise_destroy_propagates_gateway_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(502, text="gateway failed") + + clients = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + ) + + with pytest.raises(BindingDestroyError, match="502"): + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="sandbox-1", + workspace_ref="sandbox-1", + destroy_workspace=True, + ) + ) + + assert clients[0].is_closed + + +@pytest.mark.anyio +async def test_enterprise_allocation_and_home_snapshots_remain_explicitly_not_implemented() -> None: + snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="https://gateway", auth_token="secret") + bindings = EnterpriseExecutionBindingBackend(gateway_endpoint="https://gateway", auth_token="secret") + + with pytest.raises(NotImplementedError, match="Execution Binding protocol"): + _ = await snapshots.initialize( + InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1") + ) + with pytest.raises(NotImplementedError, match="Execution Binding protocol"): + _ = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), + source=cast(RuntimeLease, object()), + ) + with pytest.raises(NotImplementedError, match="Execution Binding protocol"): + await snapshots.delete("snapshot-1") + with pytest.raises(NotImplementedError, match="Execution Binding protocol"): + _ = await bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="snapshot-1", + ) + ) diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_leases.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_leases.py new file mode 100644 index 00000000000..9faecc90786 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_leases.py @@ -0,0 +1,33 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from dify_agent.runtime_backend.leases import open_runtime_lease + + +@pytest.mark.anyio +async def test_runtime_lease_body_error_survives_release_error() -> None: + lease = MagicMock() + backend = MagicMock() + backend.acquire = AsyncMock(return_value=lease) + backend.release = AsyncMock(side_effect=RuntimeError("release failed")) + + with pytest.raises(ValueError, match="body failed"): + async with open_runtime_lease(backend, "binding-ref"): + raise ValueError("body failed") + + backend.release.assert_awaited_once_with(lease) + + +@pytest.mark.anyio +async def test_runtime_lease_release_error_propagates_after_successful_body() -> None: + lease = MagicMock() + backend = MagicMock() + backend.acquire = AsyncMock(return_value=lease) + backend.release = AsyncMock(side_effect=RuntimeError("release failed")) + + with pytest.raises(RuntimeError, match="release failed"): + async with open_runtime_lease(backend, "binding-ref"): + pass + + backend.release.assert_awaited_once_with(lease) diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py new file mode 100644 index 00000000000..a1ee4d56cac --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py @@ -0,0 +1,446 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +import shlex +from typing import Mapping + +import pytest +from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView + +from dify_agent.runtime_backend import ( + BindingCreateError, + BindingDestroyError, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotCreateSpec, + HomeSnapshotCreateError, + InitializeHomeSnapshotSpec, +) +from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend + + +@dataclass(slots=True) +class _RunCall: + commands: tuple[tuple[str, ...], ...] + cwd: str | None + env: Mapping[str, str] | None + + +@dataclass(slots=True) +class _Client: + runs: list[_RunCall] + closed: bool = False + exit_code: int = 0 + output: str = "" + close_error: Exception | None = None + exit_codes: list[int] = field(default_factory=list) + + async def run( + self, + script: str, + *, + cwd: str | None = None, + env: Mapping[str, str] | None = None, + timeout: float = 10.0, + ) -> JobResult: + del timeout + commands = tuple( + tuple(shlex.split(line)) for line in script.splitlines() if line.strip() and line.strip() != "set -eu" + ) + self.runs.append(_RunCall(commands=commands, cwd=cwd, env=env)) + return JobResult( + job_id=f"job-{len(self.runs)}", + status=JobStatusName.EXITED, + done=True, + exit_code=self.exit_codes.pop(0) if self.exit_codes else self.exit_code, + output_path="/tmp/output.log", + output=self.output, + offset=0, + truncated=False, + ) + + async def wait(self, job_id: str, *, offset: int, timeout: float = 10.0) -> JobResult: + raise AssertionError((job_id, offset, timeout)) + + async def input(self, job_id: str, text: str, *, offset: int, timeout: float = 10.0) -> JobResult: + raise AssertionError((job_id, text, offset, timeout)) + + async def tail(self, job_id: str) -> JobResult: + raise AssertionError(job_id) + + async def terminate(self, job_id: str, grace_seconds: float = 2.0) -> JobStatusView: + raise AssertionError((job_id, grace_seconds)) + + async def delete( + self, + job_id: str, + *, + force: bool = False, + grace_seconds: float | None = None, + ) -> DeleteJobResponse: + del force, grace_seconds + return DeleteJobResponse(job_id=job_id) + + async def close(self) -> None: + self.closed = True + if self.close_error is not None: + raise self.close_error + + +@dataclass(slots=True) +class _Factory: + clients: list[_Client] = field(default_factory=list) + runs: list[_RunCall] = field(default_factory=list) + + def __call__(self) -> _Client: + client = _Client(runs=self.runs) + self.clients.append(client) + return client + + @property + def commands(self) -> tuple[tuple[str, ...], ...]: + return tuple(command for run in self.runs for command in run.commands) + + +@dataclass(slots=True) +class _FailingFactory: + clients: list[_Client] = field(default_factory=list) + runs: list[_RunCall] = field(default_factory=list) + + def __call__(self) -> _Client: + client = _Client( + runs=self.runs, + exit_code=1, + output="primary shellctl failure", + close_error=RuntimeError("secondary close failure"), + ) + self.clients.append(client) + return client + + +@dataclass(slots=True) +class _FailThenSucceedFactory: + clients: list[_Client] = field(default_factory=list) + runs: list[_RunCall] = field(default_factory=list) + + def __call__(self) -> _Client: + client = _Client( + runs=self.runs, + output="primary shellctl failure", + exit_codes=[1, 0], + ) + self.clients.append(client) + return client + + @property + def commands(self) -> tuple[tuple[str, ...], ...]: + return tuple(command for run in self.runs for command in run.commands) + + +@pytest.mark.anyio +async def test_local_snapshot_initialize_creates_private_snapshot_directory() -> None: + factory = _Factory() + snapshots = LocalHomeSnapshotBackend( + endpoint="http://shellctl", + auth_token="", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + snapshot_ref = await snapshots.initialize( + InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1") + ) + + assert snapshot_ref == "home-home-1" + assert ("mkdir", "-p", "/snapshots/home-home-1") in factory.commands + assert ("chmod", "700", "/snapshots/home-home-1") in factory.commands + + +@pytest.mark.anyio +async def test_local_snapshot_create_failure_removes_partial_snapshot() -> None: + factory = _FailThenSucceedFactory() + snapshots = LocalHomeSnapshotBackend( + endpoint="http://shellctl", + auth_token="", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + with pytest.raises(HomeSnapshotCreateError, match="primary shellctl failure"): + await snapshots.initialize( + InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1") + ) + + assert ("rm", "-rf", "--", "/snapshots/home-home-1") in factory.commands + + +@pytest.mark.anyio +async def test_local_binding_create_materializes_home_and_new_workspace() -> None: + factory = _Factory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + materialized_home_root="/homes", + workspace_root="/workspaces", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + allocation = await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="home-home-1", + ) + ) + + assert allocation.binding_ref == "binding-1:workspace-1" + assert allocation.workspace_ref == "workspace-1" + assert ("test", "-d", "/snapshots/home-home-1") in factory.commands + assert ("mkdir", "-p", "/workspaces/workspace-1") in factory.commands + assert ("mkdir", "-p", "/homes/binding-1") in factory.commands + assert ("cp", "-a", "/snapshots/home-home-1/.", "/homes/binding-1/") in factory.commands + assert ("chmod", "700", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands + + +@pytest.mark.anyio +async def test_local_binding_create_failure_removes_partial_home_and_workspace() -> None: + factory = _FailThenSucceedFactory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + materialized_home_root="/homes", + workspace_root="/workspaces", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + with pytest.raises(BindingCreateError, match="primary shellctl failure"): + await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="home-home-1", + ) + ) + + assert ("rm", "-rf", "--", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands + + +@pytest.mark.anyio +async def test_local_binding_acquire_scopes_commands_to_materialized_home_and_workspace() -> None: + factory = _Factory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + materialized_home_root="/homes", + workspace_root="/workspaces", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + lease = await backend.acquire("binding-1:workspace-1") + + assert lease.layout.home_dir == "/homes/binding-1" + assert lease.layout.workspace_dir == "/workspaces/workspace-1" + assert ("test", "-d", "/homes/binding-1") in factory.commands + assert ("test", "-d", "/workspaces/workspace-1") in factory.commands + + await lease.commands.run("pwd", cwd=None, env={"HOME": "/homes/other"}, timeout=10.0) + pwd_run = next(run for run in factory.runs if run.commands == (("pwd",),)) + assert pwd_run.cwd == "/workspaces/workspace-1" + assert pwd_run.env == {"HOME": "/homes/binding-1"} + with pytest.raises(ValueError, match="outside this RuntimeLease"): + await lease.commands.run("cat secret", cwd="/homes/other", timeout=10.0) + await backend.release(lease) + + +@pytest.mark.anyio +async def test_local_snapshot_checkpoint_copies_only_materialized_home() -> None: + factory = _Factory() + snapshots = LocalHomeSnapshotBackend( + endpoint="http://shellctl", + auth_token="", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + bindings = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + materialized_home_root="/homes", + workspace_root="/workspaces", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + lease = await bindings.acquire("binding-1:workspace-1") + + snapshot_ref = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), + source=lease, + ) + await bindings.release(lease) + + assert snapshot_ref == "home-home-2" + assert ("test", "-d", "/homes/binding-1") in factory.commands + assert ("mkdir", "-p", "/snapshots/home-home-2") in factory.commands + assert ("cp", "-a", "/homes/binding-1/.", "/snapshots/home-home-2/") in factory.commands + assert ("cp", "-a", "/workspaces/workspace-1/.", "/snapshots/home-home-2/") not in factory.commands + + +@pytest.mark.anyio +async def test_local_binding_destroy_removes_home_and_requested_workspace() -> None: + factory = _Factory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + materialized_home_root="/homes", + workspace_root="/workspaces", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="binding-1:workspace-1", + workspace_ref="workspace-1", + destroy_workspace=True, + ) + ) + + assert ("rm", "-rf", "--", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands + + +@pytest.mark.anyio +async def test_local_snapshot_delete_removes_snapshot_directory() -> None: + factory = _Factory() + backend = LocalHomeSnapshotBackend( + endpoint="http://shellctl", + auth_token="", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + await backend.delete("home-home-2") + + assert ("rm", "-rf", "--", "/snapshots/home-home-2") in factory.commands + + +@pytest.mark.anyio +async def test_local_backend_materializes_same_agent_twice_in_one_workspace() -> None: + factory = _Factory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + first = await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="home-home-1", + ) + ) + second = await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-2", + workspace_id="workspace-1", + existing_workspace_ref=first.workspace_ref, + home_snapshot_ref="home-home-1", + ) + ) + + assert first.binding_ref == "binding-1:workspace-1" + assert second.binding_ref == "binding-2:workspace-1" + assert first.workspace_ref == second.workspace_ref == "workspace-1" + first_lease = await backend.acquire(first.binding_ref) + second_lease = await backend.acquire(second.binding_ref) + assert first_lease.layout.home_dir != second_lease.layout.home_dir + assert first_lease.layout.workspace_dir == second_lease.layout.workspace_dir + await backend.release(first_lease) + await backend.release(second_lease) + + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=first.binding_ref, + destroy_workspace=False, + ) + ) + surviving_lease = await backend.acquire(second.binding_ref) + assert surviving_lease.layout.home_dir == "/home/dify/.dify-agent-materialized-homes/binding-2" + assert surviving_lease.layout.workspace_dir == "/home/dify/.dify-agent-workspaces/workspace-1" + await backend.release(surviving_lease) + + workspace_dir = "/home/dify/.dify-agent-workspaces/workspace-1" + assert ("test", "-d", workspace_dir) in factory.commands + assert factory.commands.count(("mkdir", "-p", workspace_dir)) == 1 + assert ( + "rm", + "-rf", + "--", + "/home/dify/.dify-agent-materialized-homes/binding-1", + ) in factory.commands + assert ( + "rm", + "-rf", + "--", + "/home/dify/.dify-agent-materialized-homes/binding-1", + workspace_dir, + ) not in factory.commands + assert ( + "cp", + "-a", + "/home/dify/.dify-agent-home-snapshots/home-home-1/.", + "/home/dify/.dify-agent-materialized-homes/binding-1/", + ) in factory.commands + assert ( + "cp", + "-a", + "/home/dify/.dify-agent-home-snapshots/home-home-1/.", + "/home/dify/.dify-agent-materialized-homes/binding-2/", + ) in factory.commands + + +@pytest.mark.anyio +async def test_local_snapshot_delete_preserves_shellctl_error_when_close_fails() -> None: + factory = _FailingFactory() + backend = LocalHomeSnapshotBackend( + endpoint="http://shellctl", + auth_token="", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + with pytest.raises(BindingDestroyError, match="primary shellctl failure"): + await backend.delete("home-1") + + assert factory.clients and all(client.closed for client in factory.clients) + + +@pytest.mark.anyio +async def test_local_binding_destroy_preserves_shellctl_error_when_close_fails() -> None: + factory = _FailingFactory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + with pytest.raises(BindingDestroyError, match="primary shellctl failure"): + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="binding-1:workspace-1", + destroy_workspace=False, + ) + ) + + assert factory.clients and all(client.closed for client in factory.clients) diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py new file mode 100644 index 00000000000..2f827be7897 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from dify_agent.runtime_backend.e2b import E2B_MAX_ACTIVE_TIMEOUT_SECONDS +from dify_agent.runtime_backend.profile import DEFAULT_E2B_TEMPLATE, RuntimeBackendSettings + + +def test_e2b_backend_uses_prepared_dify_template_by_default() -> None: + settings = RuntimeBackendSettings(runtime_backend="e2b", e2b_api_key="secret") + + assert settings.e2b_template == "difys-default-team/dify-agent-local-sandbox" + assert settings.e2b_template == DEFAULT_E2B_TEMPLATE + assert settings.e2b_active_timeout_seconds == E2B_MAX_ACTIVE_TIMEOUT_SECONDS + + +def test_e2b_backend_requires_api_key() -> None: + with pytest.raises(ValidationError, match="e2b_api_key"): + _ = RuntimeBackendSettings(runtime_backend="e2b") + + +def test_e2b_backend_rejects_active_timeout_above_platform_limit() -> None: + with pytest.raises(ValidationError, match="less than or equal"): + _ = RuntimeBackendSettings( + runtime_backend="e2b", + e2b_api_key="secret", + e2b_active_timeout_seconds=E2B_MAX_ACTIVE_TIMEOUT_SECONDS + 1, + ) + + +def test_local_backend_requires_shellctl_endpoint() -> None: + with pytest.raises(ValidationError, match="local_sandbox_endpoint"): + _ = RuntimeBackendSettings(runtime_backend="local") diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py new file mode 100644 index 00000000000..9d07ccfb027 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import cast + +import pytest + +from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellCommandStatus +from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol +from dify_agent.runtime_backend.protocols import RuntimeLayout +from dify_agent.runtime_backend.shellctl import ( + create_owned_shellctl_lease, + create_shellctl_lease, + run_shellctl_control_command, +) + + +@dataclass(slots=True) +class _FakeClient: + close_error: Exception | None = None + close_calls: int = 0 + + async def close(self) -> None: + self.close_calls += 1 + if self.close_error is not None: + raise self.close_error + + +@dataclass(slots=True) +class _FakeTransport: + close_error: Exception | None = None + close_calls: int = 0 + + async def aclose(self) -> None: + self.close_calls += 1 + if self.close_error is not None: + raise self.close_error + + +@dataclass(slots=True) +class _FakeCommands: + initial: ShellCommandResult + wait_error: Exception | None = None + delete_error: Exception | None = None + delete_calls: list[tuple[str, bool]] = field(default_factory=list) + + async def run( + self, + script: str, + *, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout: float, + ) -> ShellCommandResult: + del script, cwd, env, timeout + return self.initial + + async def wait(self, job_id: str, *, offset: int, timeout: float) -> ShellCommandResult: + del job_id, offset, timeout + if self.wait_error is not None: + raise self.wait_error + raise AssertionError("wait was not expected") + + async def read_output(self, job_id: str, *, offset: int) -> ShellCommandResult: + raise AssertionError((job_id, offset)) + + async def input(self, job_id: str, text: str, *, offset: int, timeout: float) -> ShellCommandResult: + raise AssertionError((job_id, text, offset, timeout)) + + async def interrupt(self, job_id: str, *, grace_seconds: float) -> ShellCommandStatus: + raise AssertionError((job_id, grace_seconds)) + + async def tail(self, job_id: str) -> ShellCommandResult: + raise AssertionError(job_id) + + async def delete( + self, + job_id: str, + *, + force: bool = False, + grace_seconds: float | None = None, + ) -> None: + del grace_seconds + self.delete_calls.append((job_id, force)) + if self.delete_error is not None: + raise self.delete_error + + +def _result(*, done: bool = True) -> ShellCommandResult: + return ShellCommandResult( + job_id="job-1", + status="exited" if done else "running", + done=done, + exit_code=0 if done else None, + output="ok", + offset=2, + truncated=False, + ) + + +@pytest.mark.anyio +async def test_owned_transport_is_closed_exactly_once() -> None: + client = _FakeClient() + transport = _FakeTransport() + lease = create_shellctl_lease( + handle="sandbox-1", + layout=RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace"), + entrypoint="http://shellctl", + token="secret", + client_factory=lambda: cast(ShellctlClientProtocol, cast(object, client)), + owned_transport=transport, + ) + + await lease.close() + await lease.close() + + assert client.close_calls == 1 + assert transport.close_calls == 1 + + +@pytest.mark.anyio +async def test_owned_transport_closes_when_client_close_fails_without_double_close() -> None: + client = _FakeClient(close_error=RuntimeError("client close failed")) + transport = _FakeTransport() + lease = create_shellctl_lease( + handle="sandbox-1", + layout=RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace"), + entrypoint="http://shellctl", + token="secret", + client_factory=lambda: cast(ShellctlClientProtocol, cast(object, client)), + owned_transport=transport, + ) + + with pytest.raises(RuntimeError, match="client close failed"): + await lease.close() + await lease.close() + + assert client.close_calls == 1 + assert transport.close_calls == 1 + + +@pytest.mark.anyio +async def test_owned_transport_closes_when_client_construction_fails() -> None: + transport = _FakeTransport() + + def fail_factory() -> ShellctlClientProtocol: + raise RuntimeError("client construction failed") + + with pytest.raises(RuntimeError, match="client construction failed"): + _ = await create_owned_shellctl_lease( + handle="sandbox-1", + layout=RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace"), + entrypoint="http://shellctl", + token="secret", + client_factory=fail_factory, + owned_transport=transport, + ) + + assert transport.close_calls == 1 + + +@pytest.mark.anyio +async def test_control_command_success_is_preserved_when_delete_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + commands = _FakeCommands(initial=_result(), delete_error=RuntimeError("delete failed")) + + with caplog.at_level("WARNING", logger="dify_agent.runtime_backend.shellctl"): + result = await run_shellctl_control_command(commands, "true") + + assert result.output == "ok" + assert commands.delete_calls == [("job-1", True)] + assert "delete failed" in caplog.text + + +@pytest.mark.anyio +async def test_control_command_error_is_preserved_when_delete_also_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + commands = _FakeCommands( + initial=_result(done=False), + wait_error=RuntimeError("command failed"), + delete_error=RuntimeError("delete failed"), + ) + + with caplog.at_level("WARNING", logger="dify_agent.runtime_backend.shellctl"): + with pytest.raises(RuntimeError, match="command failed"): + _ = await run_shellctl_control_command(commands, "false") + + assert commands.delete_calls == [("job-1", True)] + assert "delete failed" in caplog.text diff --git a/dify-agent/tests/local/dify_agent/server/test_app.py b/dify-agent/tests/local/dify_agent/server/test_app.py index f8f0bb57ffa..1891623b38d 100644 --- a/dify-agent/tests/local/dify_agent/server/test_app.py +++ b/dify-agent/tests/local/dify_agent/server/test_app.py @@ -8,8 +8,6 @@ import httpx import pytest from fastapi.testclient import TestClient -from dify_agent.adapters.shell.shellctl import ShellctlProvider - import dify_agent.server.app as app_module from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer @@ -17,6 +15,9 @@ from dify_agent.layers.knowledge.configs import DifyKnowledgeBaseLayerConfig from dify_agent.layers.knowledge.layer import DifyKnowledgeBaseLayer from dify_agent.layers.shell import DifyShellLayerConfig from dify_agent.layers.shell.layer import DifyShellLayer +from dify_agent.layers.runtime import DifyRuntimeLayerConfig +from dify_agent.layers.runtime.layer import DifyRuntimeLayer +from dify_agent.runtime_backend.local import LocalExecutionBindingBackend from dify_agent.runtime.compositor_factory import DifyAgentLayerProvider from dify_agent.server.app import create_app, create_dify_api_inner_http_client, create_plugin_daemon_http_client from dify_agent.server.settings import ServerSettings @@ -191,8 +192,8 @@ def test_create_app_creates_scheduler_and_closes_after_shutdown(monkeypatch: pyt plugin_daemon_api_key="daemon-secret", inner_api_url="http://dify-api", inner_api_key="inner-secret", - shellctl_entrypoint="http://shellctl", - shellctl_auth_token="shell-secret", + local_sandbox_endpoint="http://shellctl", + local_sandbox_auth_token="shell-secret", agent_stub_api_base_url="https://agent.example.com/agent-stub", server_secret_key=_base64url_secret(b"1" * 32), outbound_http_connect_timeout=1, @@ -251,7 +252,10 @@ def test_create_app_creates_scheduler_and_closes_after_shutdown(monkeypatch: pyt assert isinstance(knowledge_layer, DifyKnowledgeBaseLayer) assert knowledge_layer.inner_api_url == "http://dify-api" assert knowledge_layer.inner_api_key == "inner-secret" - assert isinstance(shell_layer.shell_provider, ShellctlProvider) + runtime_provider = next(provider for provider in layer_providers if provider.type_id == "dify.runtime") + runtime_layer = runtime_provider.create_layer(DifyRuntimeLayerConfig(backend_binding_ref="binding-1")) + assert isinstance(runtime_layer, DifyRuntimeLayer) + assert isinstance(runtime_layer.backend, LocalExecutionBindingBackend) assert shell_layer.agent_stub_api_base_url == "https://agent.example.com/agent-stub" http_client = scheduler.plugin_daemon_http_client assert http_client is fake_http_client diff --git a/dify-agent/tests/local/dify_agent/server/test_execution_bindings.py b/dify-agent/tests/local/dify_agent/server/test_execution_bindings.py new file mode 100644 index 00000000000..4ff354be6c9 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/server/test_execution_bindings.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass, field + +import pytest + +from dify_agent.protocol import CreateExecutionBindingRequest, DestroyExecutionBindingRequest +from dify_agent.runtime_backend import ( + ExecutionBindingAllocation, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, +) +from dify_agent.server.execution_bindings import ExecutionBindingService + + +@dataclass(slots=True) +class _Backend: + created: list[ExecutionBindingCreateSpec] = field(default_factory=list) + destroyed: list[ExecutionBindingDestroySpec] = field(default_factory=list) + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + self.created.append(spec) + return ExecutionBindingAllocation(binding_ref="opaque-binding", workspace_ref="opaque-workspace") + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + self.destroyed.append(spec) + + +@pytest.mark.anyio +async def test_execution_binding_service_forwards_final_contract() -> None: + backend = _Backend() + service = ExecutionBindingService(backend=backend) # pyright: ignore[reportArgumentType] + + response = await service.create_binding( + CreateExecutionBindingRequest( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="home-ref", + ) + ) + await service.destroy_binding( + DestroyExecutionBindingRequest( + binding_ref=response.binding_ref, + workspace_ref=response.workspace_ref, + destroy_workspace=True, + ) + ) + + assert backend.created == [ + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="home-ref", + ) + ] + assert backend.destroyed == [ + ExecutionBindingDestroySpec( + binding_ref="opaque-binding", + workspace_ref="opaque-workspace", + destroy_workspace=True, + ) + ] diff --git a/dify-agent/tests/local/dify_agent/server/test_home_snapshots.py b/dify-agent/tests/local/dify_agent/server/test_home_snapshots.py new file mode 100644 index 00000000000..4cee34051c3 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/server/test_home_snapshots.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import cast + +import pytest + +from dify_agent.protocol import ( + CreateHomeSnapshotFromBindingRequest, + DeleteHomeSnapshotRequest, + InitializeHomeSnapshotRequest, +) +from dify_agent.runtime_backend import HomeSnapshotCreateSpec, InitializeHomeSnapshotSpec, RuntimeLease +from dify_agent.server.home_snapshots import HomeSnapshotService + + +@dataclass(slots=True) +class _HomeBackend: + initialized: list[InitializeHomeSnapshotSpec] = field(default_factory=list) + checkpointed: list[tuple[HomeSnapshotCreateSpec, RuntimeLease]] = field(default_factory=list) + deleted: list[str] = field(default_factory=list) + + async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str: + self.initialized.append(spec) + return "snapshot-initial" + + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + self.checkpointed.append((spec, source)) + return "snapshot-build" + + async def delete(self, snapshot_ref: str) -> None: + self.deleted.append(snapshot_ref) + + +@dataclass(slots=True) +class _BindingBackend: + lease: RuntimeLease + acquired: list[str] = field(default_factory=list) + released: list[RuntimeLease] = field(default_factory=list) + + async def acquire(self, binding_ref: str) -> RuntimeLease: + self.acquired.append(binding_ref) + return self.lease + + async def release(self, lease: RuntimeLease) -> None: + self.released.append(lease) + + +@pytest.mark.anyio +async def test_home_snapshot_service_initializes_and_checkpoints_exact_binding() -> None: + lease = cast(RuntimeLease, object()) + homes = _HomeBackend() + bindings = _BindingBackend(lease=lease) + service = HomeSnapshotService( + home_snapshots=homes, # pyright: ignore[reportArgumentType] + execution_bindings=bindings, # pyright: ignore[reportArgumentType] + ) + + initial = await service.initialize( + InitializeHomeSnapshotRequest(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1") + ) + checkpoint = await service.create_from_binding( + CreateHomeSnapshotFromBindingRequest( + tenant_id="tenant-1", + agent_id="agent-1", + home_snapshot_id="home-2", + backend_binding_ref="binding-ref", + ) + ) + + assert initial.snapshot_ref == "snapshot-initial" + assert checkpoint.snapshot_ref == "snapshot-build" + assert bindings.acquired == ["binding-ref"] + assert bindings.released == [lease] + assert homes.checkpointed == [ + ( + HomeSnapshotCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + home_snapshot_id="home-2", + ), + lease, + ) + ] + + await service.delete(DeleteHomeSnapshotRequest(snapshot_ref="snapshot-build")) + assert homes.deleted == ["snapshot-build"] + + +@pytest.mark.anyio +async def test_snapshot_checkpoint_releases_binding_when_create_fails() -> None: + lease = cast(RuntimeLease, object()) + + class _FailingHomeBackend(_HomeBackend): + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + del spec, source + raise RuntimeError("checkpoint failed") + + homes = _FailingHomeBackend() + bindings = _BindingBackend(lease=lease) + service = HomeSnapshotService( + home_snapshots=homes, # pyright: ignore[reportArgumentType] + execution_bindings=bindings, # pyright: ignore[reportArgumentType] + ) + + with pytest.raises(RuntimeError, match="checkpoint failed"): + await service.create_from_binding( + CreateHomeSnapshotFromBindingRequest( + tenant_id="tenant-1", + agent_id="agent-1", + home_snapshot_id="home-2", + backend_binding_ref="binding-ref", + ) + ) + + assert bindings.released == [lease] diff --git a/dify-agent/tests/local/dify_agent/server/test_sandbox_files.py b/dify-agent/tests/local/dify_agent/server/test_sandbox_files.py deleted file mode 100644 index 6bbe954ad7f..00000000000 --- a/dify-agent/tests/local/dify_agent/server/test_sandbox_files.py +++ /dev/null @@ -1,628 +0,0 @@ -from __future__ import annotations - -import asyncio -import base64 -import json -import os -from pathlib import Path -import subprocess -import sys -import types -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from typing import Literal, cast - -import pytest -from agenton.compositor import CompositorSessionSnapshot, LayerProvider -from agenton.compositor.schemas import LayerSessionSnapshot -from agenton.layers.base import LifecycleState -from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol, ShellctlProvider -from dify_agent.agent_stub.shell_env import ( - AGENT_STUB_API_BASE_URL_ENV_VAR, - AGENT_STUB_AUTH_JWE_ENV_VAR, - AGENT_STUB_DRIVE_BASE_ENV_VAR, -) - -if "graphon.model_runtime.entities.llm_entities" not in sys.modules: - graphon_module = types.ModuleType("graphon") - model_runtime_module = types.ModuleType("graphon.model_runtime") - entities_module = types.ModuleType("graphon.model_runtime.entities") - llm_entities_module = types.ModuleType("graphon.model_runtime.entities.llm_entities") - message_entities_module = types.ModuleType("graphon.model_runtime.entities.message_entities") - - llm_entities_module.LLMResultChunk = type("LLMResultChunk", (), {}) - llm_entities_module.LLMUsage = type("LLMUsage", (), {}) - - for name in ( - "AssistantPromptMessage", - "AudioPromptMessageContent", - "DocumentPromptMessageContent", - "ImagePromptMessageContent", - "PromptMessage", - "PromptMessageContentUnionTypes", - "PromptMessageTool", - "SystemPromptMessage", - "TextPromptMessageContent", - "ToolPromptMessage", - "UserPromptMessage", - "VideoPromptMessageContent", - ): - setattr(message_entities_module, name, type(name, (), {})) - - sys.modules["graphon"] = graphon_module - sys.modules["graphon.model_runtime"] = model_runtime_module - sys.modules["graphon.model_runtime.entities"] = entities_module - sys.modules["graphon.model_runtime.entities.llm_entities"] = llm_entities_module - sys.modules["graphon.model_runtime.entities.message_entities"] = message_entities_module - - graphon_module.model_runtime = model_runtime_module - model_runtime_module.entities = entities_module - entities_module.llm_entities = llm_entities_module - entities_module.message_entities = message_entities_module - -if "jsonschema" not in sys.modules: - jsonschema_module = types.ModuleType("jsonschema") - jsonschema_exceptions_module = types.ModuleType("jsonschema.exceptions") - jsonschema_protocols_module = types.ModuleType("jsonschema.protocols") - jsonschema_validators_module = types.ModuleType("jsonschema.validators") - - class _SchemaError(Exception): - pass - - class _ValidationError(Exception): - path: tuple[object, ...] = () - - class _Validator: - @staticmethod - def check_schema(schema): - return None - - def __init__(self, schema): - self.schema = schema - - def iter_errors(self, value): - return iter(()) - - def _validator_for(schema): - return _Validator - - jsonschema_module.SchemaError = _SchemaError - jsonschema_exceptions_module.ValidationError = _ValidationError - jsonschema_protocols_module.Validator = _Validator - jsonschema_validators_module.validator_for = _validator_for - - sys.modules["jsonschema"] = jsonschema_module - sys.modules["jsonschema.exceptions"] = jsonschema_exceptions_module - sys.modules["jsonschema.protocols"] = jsonschema_protocols_module - sys.modules["jsonschema.validators"] = jsonschema_validators_module - -from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig -from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer -from dify_agent.layers.shell import DifyShellLayerConfig -from dify_agent.layers.shell.layer import CompleteRemoteCommandResult, DifyShellLayer -from dify_agent.protocol import ( - CreateRunRequest, - RunComposition, - RunLayerSpec, - SandboxListRequest, - SandboxLocator, - SandboxReadRequest, - SandboxUploadRequest, - build_sandbox_locator_from_run_request, -) -from dify_agent.server.sandbox_files import ( - _LIST_SCRIPT, - SandboxFileError, - SandboxFileService, - _OUTPUT_BEGIN, - _OUTPUT_END, - _READ_SCRIPT, - _UPLOAD_SCRIPT, - _decode_sandbox_payload, - _shell_result_details, -) - - -@dataclass(slots=True) -class _Job: - job_id: str - status: str = "exited" - done: bool = True - exit_code: int | None = 0 - output: str = "" - offset: int = 0 - truncated: bool = False - output_path: str | None = "/tmp/sandbox-job.out" - - -@dataclass(slots=True) -class RunCall: - script: str - cwd: str | None - env: dict[str, str] | None - timeout: float - - -class FakeShellctlClient: - def __init__(self, *, run_handler: Callable[[str, str | None, dict[str, str] | None, float], _Job]) -> None: - self.run_handler = run_handler - self.run_calls: list[RunCall] = [] - self.delete_calls: list[str] = [] - - async def run( - self, script: str, *, cwd: str | None = None, env: dict[str, str] | None = None, timeout: float = 10.0 - ) -> _Job: - self.run_calls.append(RunCall(script=script, cwd=cwd, env=env, timeout=timeout)) - return self.run_handler(script, cwd, env, timeout) - - async def wait(self, job_id: str, *, offset: int, timeout: float = 10.0) -> _Job: - raise AssertionError(f"Unexpected wait() call for {job_id} offset={offset} timeout={timeout}") - - async def input(self, job_id: str, text: str, *, offset: int, timeout: float = 10.0) -> _Job: - raise AssertionError(f"Unexpected input() call for {job_id} text={text!r}") - - async def tail(self, job_id: str) -> _Job: - raise AssertionError(f"Unexpected tail() call for {job_id}") - - async def terminate(self, job_id: str, grace_seconds: float = 10.0) -> _Job: - raise AssertionError(f"Unexpected terminate() call for {job_id} grace={grace_seconds}") - - async def delete(self, job_id: str, *, force: bool = False, grace_seconds: float | None = None) -> None: - del force, grace_seconds - self.delete_calls.append(job_id) - return None - - async def close(self) -> None: - return None - - -def _wrap(payload: dict[str, object], *, pty_wrap: int = 0, noise: bool = False) -> str: - blob = base64.b64encode(json.dumps(payload).encode("utf-8")).decode("ascii") - if pty_wrap: - blob = "\n".join(blob[index : index + pty_wrap] for index in range(0, len(blob), pty_wrap)) - framed = f"{_OUTPUT_BEGIN}{blob}{_OUTPUT_END}\n" - if noise: - framed = f"user@host$ python3 - ...\r\n{framed}user@host$ \r\n" - return framed - - -def _complete_result( - *, - output: str, - exit_code: int | None = 0, - output_complete: bool = True, - incomplete_reason: Literal["output_limit", "timeout"] | None = None, - job_id: str = "sandbox-job", -) -> CompleteRemoteCommandResult: - return CompleteRemoteCommandResult( - job_id=job_id, - status="exited", - done=True, - exit_code=exit_code, - output=output, - output_complete=output_complete, - incomplete_reason=incomplete_reason, - offset=len(output), - output_path="/tmp/sandbox-job.out", - ) - - -def _run_embedded_script( - script_source: str, - *, - args: list[str], - cwd: Path, - env: Mapping[str, str] | None = None, -) -> dict[str, object]: - merged_env = dict(os.environ) - if env is not None: - merged_env.update(env) - completed = subprocess.run( - [sys.executable, "-", *args], - input=script_source, - text=True, - capture_output=True, - cwd=cwd, - env=merged_env, - check=False, - ) - return _decode_sandbox_payload(_complete_result(output=completed.stdout, exit_code=completed.returncode)) - - -def _execution_context() -> DifyExecutionContextLayerConfig: - return DifyExecutionContextLayerConfig( - tenant_id="tenant-1", - user_id="user-1", - user_from="account", - app_id="app-1", - conversation_id="conv-1", - agent_id="agent-1", - agent_config_version_id="snapshot-1", - agent_mode="agent_app", - invoke_from="service-api", - ) - - -def _locator() -> SandboxLocator: - request = CreateRunRequest( - composition=RunComposition( - layers=[ - RunLayerSpec(name="execution_context", type="dify.execution_context", config=_execution_context()), - RunLayerSpec( - name="shell", - type="dify.shell", - deps={"execution_context": "execution_context"}, - config=DifyShellLayerConfig(agent_stub_drive_ref="agent-1"), - ), - ] - ), - session_snapshot=CompositorSessionSnapshot( - layers=[ - LayerSessionSnapshot( - name="execution_context", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={} - ), - LayerSessionSnapshot( - name="shell", - lifecycle_state=LifecycleState.SUSPENDED, - runtime_state={"session_id": "abc12ff", "workspace_cwd": "~/workspace/abc12ff"}, - ), - ] - ), - ) - return build_sandbox_locator_from_run_request(request) - - -def _service( - run_handler: Callable[[str, str | None, dict[str, str] | None, float], _Job], -) -> tuple[SandboxFileService, FakeShellctlClient]: - client = FakeShellctlClient(run_handler=run_handler) - execution_context_provider = LayerProvider.from_factory( - layer_type=DifyExecutionContextLayer, - create=lambda config: DifyExecutionContextLayer.from_config_with_settings( - DifyExecutionContextLayerConfig.model_validate(config), - daemon_url="http://plugin-daemon", - daemon_api_key="daemon-secret", - ), - ) - shell_provider = LayerProvider.from_factory( - layer_type=DifyShellLayer, - create=lambda config: DifyShellLayer.from_config_with_settings( - DifyShellLayerConfig.model_validate(config), - shell_provider=ShellctlProvider( - entrypoint="http://shellctl", - token="", - client_factory=lambda: cast(ShellctlClientProtocol, cast(object, client)), - ), - agent_stub_api_base_url="https://agent.example.com/agent-stub", - agent_stub_token_factory=lambda execution_context, *, session_id: ( - f"token-for:{execution_context.tenant_id}:{session_id}" - ), - ), - ) - return SandboxFileService(layer_providers=(execution_context_provider, shell_provider)), client - - -def _sandbox_python_run_call(client: FakeShellctlClient) -> RunCall: - for run_call in reversed(client.run_calls): - if run_call.script.startswith("python3 - "): - return run_call - raise AssertionError("sandbox python script was not executed") - - -def _sandbox_list_entries(payload: dict[str, object]) -> list[object]: - entries = payload.get("entries") - assert isinstance(entries, list) - return entries - - -def test_list_files_runs_fixed_script_and_parses_response() -> None: - service, client = _service( - lambda script, cwd, env, timeout: _Job( - job_id="sandbox-job", - output=_wrap( - { - "path": ".", - "entries": [{"name": "notes.txt", "type": "file", "size": 5, "mtime": 1}], - "truncated": False, - } - ), - ) - ) - - result = asyncio.run(service.list_files(SandboxListRequest(locator=_locator(), path="."))) - - assert result.entries[0].name == "notes.txt" - script_call = _sandbox_python_run_call(client) - assert script_call.cwd == "/home/agent-1/workspace/abc12ff" - assert script_call.env == {"HOME": "/home/agent-1"} - assert "python3 - . 1000 <<'PY'" in script_call.script - assert client.delete_calls[-1] == "sandbox-job" - - -def test_list_files_allows_parent_relative_paths() -> None: - service, client = _service( - lambda script, cwd, env, timeout: _Job( - job_id="sandbox-job", - output=_wrap({"path": "../shared", "entries": [], "truncated": False}), - ) - ) - - result = asyncio.run(service.list_files(SandboxListRequest(locator=_locator(), path="../shared"))) - - assert result.path == "../shared" - assert "python3 - ../shared 1000 <<'PY'" in _sandbox_python_run_call(client).script - - -@pytest.mark.parametrize( - ("path", "expected_command"), - [ - ("~", "python3 - '~' 1000 <<'PY'"), - ("~/shared", "python3 - '~/shared' 1000 <<'PY'"), - ], -) -def test_list_files_allows_home_relative_paths(path: str, expected_command: str) -> None: - service, client = _service( - lambda script, cwd, env, timeout: _Job( - job_id="sandbox-job", - output=_wrap({"path": path, "entries": [], "truncated": False}), - ) - ) - - result = asyncio.run(service.list_files(SandboxListRequest(locator=_locator(), path=path))) - - assert result.path == path - assert expected_command in _sandbox_python_run_call(client).script - - -def test_embedded_scripts_allow_parent_relative_paths(tmp_path: Path) -> None: - workspace_dir = tmp_path / "workspace" - cwd = workspace_dir / "run" - shared_dir = workspace_dir / "shared" - cwd.mkdir(parents=True) - shared_dir.mkdir() - notes_path = shared_dir / "notes.txt" - notes_path.write_text("hello", encoding="utf-8") - - bin_dir = tmp_path / "bin" - bin_dir.mkdir() - fake_dify_agent = bin_dir / "dify-agent" - fake_dify_agent.write_text( - "\n".join( - [ - "#!/usr/bin/env python3", - "import json", - "import sys", - 'if sys.argv[1:] != ["file", "upload", "../shared/notes.txt"]:', - ' raise SystemExit(f"unexpected args: {sys.argv[1:]!r}")', - 'print(json.dumps({"transfer_method": "tool_file", "reference": "file-ref", "download_url": "https://files.example.com/notes.txt"}))', - ] - ) - + "\n", - encoding="utf-8", - ) - fake_dify_agent.chmod(0o755) - - list_payload = _run_embedded_script(_LIST_SCRIPT, args=["../shared", "1000"], cwd=cwd) - read_payload = _run_embedded_script(_READ_SCRIPT, args=["../shared/notes.txt", "8"], cwd=cwd) - upload_payload = _run_embedded_script( - _UPLOAD_SCRIPT, - args=["../shared/notes.txt"], - cwd=cwd, - env={"PATH": f"{bin_dir}{os.pathsep}{os.environ.get('PATH', os.defpath)}"}, - ) - - assert list_payload["path"] == "../shared" - assert any( - isinstance(entry, dict) and entry.get("name") == "notes.txt" for entry in _sandbox_list_entries(list_payload) - ) - assert read_payload == { - "path": "../shared/notes.txt", - "size": 5, - "truncated": False, - "binary": False, - "text": "hello", - } - assert upload_payload == { - "path": "../shared/notes.txt", - "file": { - "transfer_method": "tool_file", - "reference": "file-ref", - "download_url": "https://files.example.com/notes.txt", - }, - } - - -def test_embedded_scripts_expand_home_relative_paths(tmp_path: Path) -> None: - cwd = tmp_path / "workspace" / "run" - cwd.mkdir(parents=True) - home_dir = tmp_path / "home" / "agent-1" - shared_dir = home_dir / "shared" - shared_dir.mkdir(parents=True) - notes_path = shared_dir / "notes.txt" - notes_path.write_text("hello", encoding="utf-8") - - bin_dir = tmp_path / "bin" - bin_dir.mkdir() - fake_dify_agent = bin_dir / "dify-agent" - fake_dify_agent.write_text( - "\n".join( - [ - "#!/usr/bin/env python3", - "import json", - "import sys", - 'if sys.argv[1:] != ["file", "upload", "~/shared/notes.txt"]:', - ' raise SystemExit(f"unexpected args: {sys.argv[1:]!r}")', - 'print(json.dumps({"transfer_method": "tool_file", "reference": "file-ref", "download_url": "https://files.example.com/notes.txt"}))', - ] - ) - + "\n", - encoding="utf-8", - ) - fake_dify_agent.chmod(0o755) - - script_env = {"HOME": str(home_dir), "PATH": f"{bin_dir}{os.pathsep}{os.environ.get('PATH', os.defpath)}"} - list_payload = _run_embedded_script(_LIST_SCRIPT, args=["~/shared", "1000"], cwd=cwd, env=script_env) - read_payload = _run_embedded_script(_READ_SCRIPT, args=["~/shared/notes.txt", "8"], cwd=cwd, env=script_env) - upload_payload = _run_embedded_script(_UPLOAD_SCRIPT, args=["~/shared/notes.txt"], cwd=cwd, env=script_env) - - assert list_payload["path"] == "~/shared" - assert any( - isinstance(entry, dict) and entry.get("name") == "notes.txt" for entry in _sandbox_list_entries(list_payload) - ) - assert read_payload == { - "path": "~/shared/notes.txt", - "size": 5, - "truncated": False, - "binary": False, - "text": "hello", - } - assert upload_payload == { - "path": "~/shared/notes.txt", - "file": { - "transfer_method": "tool_file", - "reference": "file-ref", - "download_url": "https://files.example.com/notes.txt", - }, - } - - -@pytest.mark.parametrize("bad_path", ["/etc/passwd", "~other/secret-dir", "bad\x00path"]) -def test_list_files_rejects_invalid_paths_before_shell_execution(bad_path: str) -> None: - service, client = _service(lambda script, cwd, env, timeout: _Job(job_id="sandbox-job", output="unused")) - - with pytest.raises(SandboxFileError, match="path"): - asyncio.run(service.list_files(SandboxListRequest(locator=_locator(), path=bad_path))) - - assert client.run_calls == [] - - -def test_decode_payload_reports_incomplete_capture_when_frame_is_missing() -> None: - with pytest.raises(SandboxFileError, match="incomplete before framed payload was captured"): - _decode_sandbox_payload( - _complete_result(output="partial", output_complete=False, incomplete_reason="output_limit") - ) - - -def test_decode_payload_reports_incomplete_capture_when_frame_is_corrupt() -> None: - broken = f"{_OUTPUT_BEGIN}%%%%{_OUTPUT_END}" - with pytest.raises(SandboxFileError, match="incomplete while decoding framed payload"): - _decode_sandbox_payload(_complete_result(output=broken, output_complete=False, incomplete_reason="timeout")) - - -def test_upload_injects_agent_stub_env_and_returns_mapping() -> None: - service, client = _service( - lambda script, cwd, env, timeout: _Job( - job_id="sandbox-job", - output=_wrap( - { - "path": "report.txt", - "file": { - "transfer_method": "tool_file", - "reference": "file-ref", - "download_url": "https://files.example.com/report.txt", - }, - }, - noise=True, - ), - ) - ) - - result = asyncio.run(service.upload_file(SandboxUploadRequest(locator=_locator(), path="report.txt"))) - - assert result.file.transfer_method == "tool_file" - assert result.file.reference == "file-ref" - assert result.file.download_url == "https://files.example.com/report.txt" - script_call = _sandbox_python_run_call(client) - assert script_call.cwd == "/home/agent-1/workspace/abc12ff" - assert script_call.env == { - "HOME": "/home/agent-1", - AGENT_STUB_API_BASE_URL_ENV_VAR: "https://agent.example.com/agent-stub", - AGENT_STUB_AUTH_JWE_ENV_VAR: "token-for:tenant-1:abc12ff", - AGENT_STUB_DRIVE_BASE_ENV_VAR: "/mnt/drive/agent-1", - } - - -def test_upload_rejects_missing_download_url_in_shell_payload() -> None: - service, _client = _service( - lambda script, cwd, env, timeout: _Job( - job_id="sandbox-job", - output=_wrap( - { - "path": "report.txt", - "file": { - "transfer_method": "tool_file", - "reference": "file-ref", - }, - } - ), - ) - ) - - with pytest.raises(SandboxFileError, match="sandbox command returned invalid payload"): - _ = asyncio.run(service.upload_file(SandboxUploadRequest(locator=_locator(), path="report.txt"))) - - -def test_shell_result_details_include_output_metadata_and_tail() -> None: - details = _shell_result_details( - _complete_result(output="hello", output_complete=False, incomplete_reason="output_limit") - ) - assert "output_complete=False" in details - assert "incomplete_reason=output_limit" in details - assert "output_path=/tmp/sandbox-job.out" in details - assert details.endswith("hello") - - -def test_read_file_uses_complete_mode_and_parses_response() -> None: - service, client = _service( - lambda script, cwd, env, timeout: _Job( - job_id="sandbox-job", - output=_wrap({"path": "notes.txt", "size": 5, "truncated": False, "binary": False, "text": "hello"}), - ) - ) - - result = asyncio.run(service.read_file(SandboxReadRequest(locator=_locator(), path="notes.txt", max_bytes=8))) - - assert result.text == "hello" - assert "python3 - notes.txt 8 <<'PY'" in _sandbox_python_run_call(client).script - - -@pytest.mark.parametrize( - ("sandbox_request", "expected_command"), - [ - (SandboxReadRequest(locator=_locator(), path="../notes.txt", max_bytes=8), "python3 - ../notes.txt 8 <<'PY'"), - (SandboxUploadRequest(locator=_locator(), path="../report.txt"), "python3 - ../report.txt <<'PY'"), - (SandboxReadRequest(locator=_locator(), path="~/notes.txt", max_bytes=8), "python3 - '~/notes.txt' 8 <<'PY'"), - (SandboxUploadRequest(locator=_locator(), path="~/report.txt"), "python3 - '~/report.txt' <<'PY'"), - ], -) -def test_read_and_upload_allow_relative_paths( - sandbox_request: SandboxReadRequest | SandboxUploadRequest, expected_command: str -) -> None: - expected_path = sandbox_request.path - service, client = _service( - lambda script, cwd, env, timeout: _Job( - job_id="sandbox-job", - output=_wrap( - {"path": expected_path, "size": 5, "truncated": False, "binary": False, "text": "hello"} - if isinstance(sandbox_request, SandboxReadRequest) - else { - "path": expected_path, - "file": { - "transfer_method": "tool_file", - "reference": "file-ref", - "download_url": "https://files.example.com/report.txt", - }, - } - ), - ) - ) - - if isinstance(sandbox_request, SandboxReadRequest): - result = asyncio.run(service.read_file(sandbox_request)) - assert result.path == expected_path - else: - result = asyncio.run(service.upload_file(sandbox_request)) - assert result.path == expected_path - assert result.file.download_url == "https://files.example.com/report.txt" - - assert expected_command in _sandbox_python_run_call(client).script diff --git a/dify-agent/tests/local/dify_agent/server/test_settings.py b/dify-agent/tests/local/dify_agent/server/test_settings.py index 25ae693963b..7d7962f55ba 100644 --- a/dify-agent/tests/local/dify_agent/server/test_settings.py +++ b/dify-agent/tests/local/dify_agent/server/test_settings.py @@ -8,12 +8,13 @@ import httpx import pytest from pydantic import ValidationError -from dify_agent.adapters.shell.enterprise import EnterpriseShellProvider -from dify_agent.adapters.shell.shellctl import ShellctlProvider from dify_agent.agent_stub.server.agent_stub_drive import DifyApiAgentStubDriveRequestHandler from dify_agent.agent_stub.server.agent_stub_files import DifyApiAgentStubFileRequestHandler from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec from dify_agent.server.settings import ServerSettings +from dify_agent.runtime_backend.e2b import E2BExecutionBindingBackend +from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend +from dify_agent.runtime_backend.local import LocalExecutionBindingBackend def _base64url_secret(value: bytes) -> str: @@ -27,7 +28,7 @@ def test_server_settings_reads_shellctl_entrypoint_from_env(monkeypatch: pytest. settings = ServerSettings() - assert settings.shellctl_entrypoint == "http://shellctl.example" + assert settings.local_sandbox_endpoint == "http://shellctl.example" def test_server_settings_reads_shellctl_auth_token_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -35,22 +36,7 @@ def test_server_settings_reads_shellctl_auth_token_from_env(monkeypatch: pytest. settings = ServerSettings() - assert settings.shellctl_auth_token == "shell-secret" - - -def test_server_settings_reads_shell_home_root_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("DIFY_AGENT_SHELL_HOME_ROOT", "/tmp/dify-agent-home/") - - settings = ServerSettings() - - assert settings.shell_home_root == "/tmp/dify-agent-home" - - -def test_server_settings_rejects_relative_shell_home_root(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("DIFY_AGENT_SHELL_HOME_ROOT", "relative/path") - - with pytest.raises(ValidationError, match="DIFY_AGENT_SHELL_HOME_ROOT must be an absolute path"): - ServerSettings() + assert settings.local_sandbox_auth_token == "shell-secret" def test_server_settings_reads_enterprise_timeouts_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -63,6 +49,14 @@ def test_server_settings_reads_enterprise_timeouts_from_env(monkeypatch: pytest. assert settings.enterprise_sandbox_proxy_timeout == 90 +def test_server_settings_reads_e2b_active_timeout_from_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS", "900") + + settings = ServerSettings() + + assert settings.e2b_active_timeout_seconds == 900 + + def test_server_settings_defaults_shellctl_auth_token_to_none( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -72,7 +66,7 @@ def test_server_settings_defaults_shellctl_auth_token_to_none( settings = ServerSettings() - assert settings.shellctl_auth_token is None + assert settings.local_sandbox_auth_token is None def test_server_settings_reads_agent_stub_settings_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -251,61 +245,84 @@ def test_server_settings_create_agent_stub_drive_request_handler_returns_handler assert timeout.pool == 44 -def test_build_shell_provider_returns_none_when_shellctl_entrypoint_is_unset( +def test_build_runtime_backend_profile_returns_none_when_local_endpoint_is_unset( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: monkeypatch.delenv("DIFY_AGENT_SHELLCTL_ENTRYPOINT", raising=False) monkeypatch.chdir(tmp_path) - assert ServerSettings().build_shell_provider() is None + assert ServerSettings().build_runtime_backend_profile() is None -def test_build_shell_provider_returns_shellctl_provider_when_configured() -> None: +def test_build_runtime_backend_profile_returns_local_drivers_when_configured() -> None: settings = ServerSettings( - shell_provider="shellctl", - shellctl_entrypoint="http://shellctl.example", - shellctl_auth_token="shell-secret", + runtime_backend="local", + local_sandbox_endpoint="http://shellctl.example", + local_sandbox_auth_token="shell-secret", ) - provider = settings.build_shell_provider() + profile = settings.build_runtime_backend_profile() - assert isinstance(provider, ShellctlProvider) - assert provider.entrypoint == "http://shellctl.example" - assert provider.token == "shell-secret" + assert profile is not None + assert isinstance(profile.execution_bindings, LocalExecutionBindingBackend) + assert profile.execution_bindings.endpoint == "http://shellctl.example" + assert profile.execution_bindings.auth_token == "shell-secret" -def test_build_shell_provider_returns_enterprise_provider_when_selected() -> None: +def test_build_runtime_backend_profile_returns_enterprise_drivers_when_selected() -> None: settings = ServerSettings( - shell_provider="enterprise", + runtime_backend="enterprise", enterprise_sandbox_gateway_endpoint="https://gateway.example", enterprise_sandbox_gateway_auth_token="gateway-secret", enterprise_sandbox_gateway_timeout=45, enterprise_sandbox_proxy_timeout=90, ) - provider = settings.build_shell_provider() + profile = settings.build_runtime_backend_profile() - assert isinstance(provider, EnterpriseShellProvider) - assert provider.gateway_endpoint == "https://gateway.example" - assert provider.auth_token == "gateway-secret" - assert provider.gateway_timeout == 45 - assert provider.proxy_timeout == 90 + assert profile is not None + assert isinstance(profile.execution_bindings, EnterpriseExecutionBindingBackend) + assert profile.execution_bindings.gateway_endpoint == "https://gateway.example" + assert profile.execution_bindings.auth_token == "gateway-secret" + assert profile.execution_bindings.gateway_timeout == 45 + assert profile.execution_bindings.proxy_timeout == 90 -def test_build_shell_provider_returns_none_when_enterprise_endpoint_is_unset( +def test_build_runtime_backend_profile_passes_e2b_active_timeout() -> None: + settings = ServerSettings( + runtime_backend="e2b", + e2b_api_key="e2b-secret", + e2b_active_timeout_seconds=900, + ) + + profile = settings.build_runtime_backend_profile() + + assert profile is not None + assert isinstance(profile.execution_bindings, E2BExecutionBindingBackend) + assert profile.execution_bindings.active_timeout_seconds == 900 + + +def test_sandbox_file_upload_limit_defaults_to_tool_file_limit() -> None: + settings = ServerSettings() + + assert settings.sandbox_file_upload_max_bytes == 50 * 1024 * 1024 + + +def test_build_runtime_backend_profile_rejects_missing_enterprise_endpoint( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: monkeypatch.delenv("DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_ENDPOINT", raising=False) monkeypatch.chdir(tmp_path) - assert ServerSettings(shell_provider="enterprise").build_shell_provider() is None + with pytest.raises(ValidationError, match="enterprise_sandbox_gateway_endpoint is required"): + _ = ServerSettings(runtime_backend="enterprise").build_runtime_backend_profile() -def test_build_shell_provider_rejects_blank_shellctl_entrypoint() -> None: - with pytest.raises(ValidationError, match="shellctl_entrypoint is required"): - _ = ServerSettings(shell_provider="shellctl", shellctl_entrypoint=" ").build_shell_provider() +def test_build_runtime_backend_profile_rejects_blank_local_endpoint() -> None: + with pytest.raises(ValidationError, match="local_sandbox_endpoint is required"): + _ = ServerSettings(runtime_backend="local", local_sandbox_endpoint=" ").build_runtime_backend_profile() def test_server_settings_parses_shell_redact_patterns_json_array(monkeypatch: pytest.MonkeyPatch) -> None: @@ -342,4 +359,4 @@ def test_server_settings_rejects_non_array_shell_redact_patterns(monkeypatch: py settings = ServerSettings() with pytest.raises(ValueError, match="must be a JSON array"): - settings.get_shell_redact_patterns() + _ = settings.get_shell_redact_patterns() diff --git a/dify-agent/tests/local/dify_agent/server/test_workspace_files.py b/dify-agent/tests/local/dify_agent/server/test_workspace_files.py new file mode 100644 index 00000000000..78af3bee9ed --- /dev/null +++ b/dify-agent/tests/local/dify_agent/server/test_workspace_files.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import cast + +import pytest + +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig +from dify_agent.protocol import ( + WorkspaceListRequest, + WorkspaceReadRequest, + WorkspaceUploadRequest, + WorkspaceUploadedFile, +) +from dify_agent.runtime_backend import ( + RuntimeLayout, + RuntimeLease, + WorkspaceFileContent, + WorkspaceFileEntry, + WorkspaceListResult, + WorkspaceReadResult, +) +from dify_agent.server.workspace_files import WorkspaceFileService + + +@dataclass(slots=True) +class _Files: + calls: list[tuple[str, str]] = field(default_factory=list) + + async def list_directory(self, *, path: str, limit: int) -> WorkspaceListResult: + self.calls.append(("list", path)) + assert limit == 1000 + return WorkspaceListResult( + path=path, + entries=(WorkspaceFileEntry(name="note.txt", type="file", size=4, mtime=1),), + truncated=False, + ) + + async def read_file(self, *, path: str, max_bytes: int) -> WorkspaceReadResult: + self.calls.append(("read", path)) + return WorkspaceReadResult(path=path, size=4, truncated=False, binary=False, text="note") + + async def read_bytes(self, *, path: str, max_bytes: int) -> WorkspaceFileContent: + self.calls.append(("bytes", path)) + return WorkspaceFileContent(path=path, size=4, content=b"note") + + +@dataclass(slots=True) +class _Lease: + files: _Files = field(default_factory=_Files) + layout: RuntimeLayout = RuntimeLayout(home_dir="/home/agent", workspace_dir="/workspace") + commands: object = field(default_factory=object) + + +@dataclass(slots=True) +class _Backend: + lease: RuntimeLease + acquired: list[str] = field(default_factory=list) + releases: int = 0 + + async def acquire(self, binding_ref: str) -> RuntimeLease: + self.acquired.append(binding_ref) + return self.lease + + async def release(self, lease: RuntimeLease) -> None: + assert lease is self.lease + self.releases += 1 + + +@dataclass(slots=True) +class _Uploader: + uploads: list[tuple[str, str, bytes]] = field(default_factory=list) + + async def upload( + self, + *, + execution_context: DifyExecutionContextLayerConfig, + filename: str, + mimetype: str, + content: bytes, + ) -> WorkspaceUploadedFile: + del execution_context + self.uploads.append((filename, mimetype, content)) + return WorkspaceUploadedFile(reference="tool-file-1", download_url="https://files/note.txt") + + +@pytest.mark.anyio +async def test_workspace_service_passes_paths_directly_and_leases_each_operation() -> None: + lease = _Lease() + backend = _Backend(lease=cast(RuntimeLease, lease)) + uploader = _Uploader() + service = WorkspaceFileService( + execution_bindings=backend, # pyright: ignore[reportArgumentType] + upload_max_bytes=1024, + file_uploader=uploader, + ) + + listing = await service.list_files(WorkspaceListRequest(backend_binding_ref="binding-ref", path="/var/data")) + preview = await service.read_file(WorkspaceReadRequest(backend_binding_ref="binding-ref", path="~/note.txt")) + uploaded = await service.upload_file( + WorkspaceUploadRequest( + backend_binding_ref="binding-ref", + path="../outside.txt", + execution_context=DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_from="account", + agent_mode="agent_app", + invoke_from="debugger", + ), + ) + ) + + assert listing.path == "/var/data" + assert preview.path == "~/note.txt" + assert uploaded.path == "../outside.txt" + assert lease.files.calls == [ + ("list", "/var/data"), + ("read", "~/note.txt"), + ("bytes", "../outside.txt"), + ] + assert backend.acquired == ["binding-ref", "binding-ref", "binding-ref"] + assert backend.releases == 3 + assert uploader.uploads == [("outside.txt", "text/plain", b"note")] diff --git a/dify-agent/tests/local/test_packaging.py b/dify-agent/tests/local/test_packaging.py index b11bcb650a8..d82c3f192b2 100644 --- a/dify-agent/tests/local/test_packaging.py +++ b/dify-agent/tests/local/test_packaging.py @@ -15,6 +15,7 @@ CLIENT_SHARED_DTO_DEPENDENCIES = { } SERVER_RUNTIME_DEPENDENCIES = { + "e2b>=2.34.0,<3.0.0", "fastapi==0.136.0", "graphon==0.5.2", "jsonschema>=4.23.0,<5.0.0", diff --git a/dify-agent/uv.lock b/dify-agent/uv.lock index 8c2ee1c876f..727a29e41f4 100644 --- a/dify-agent/uv.lock +++ b/dify-agent/uv.lock @@ -198,6 +198,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/de/acae8e9f9a1f4bb393d41c8265898b0f29772e38eac14e9f69d191e2c006/blis-1.3.3-cp314-cp314-win_amd64.whl", hash = "sha256:9e5fdf4211b1972400f8ff6dafe87cb689c5d84f046b4a76b207c0bd2270faaf", size = 6324695, upload-time = "2025-11-17T12:28:28.401Z" }, ] +[[package]] +name = "bracex" +version = "3.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ac/01/5f394b8bcd6e5b92f73130990960423bbb19711f906bd9fe9ea5557c667c/bracex-3.0.1.tar.gz", hash = "sha256:4e38e32392e4a4780fe15d644bfc7c8514057cfc3861e060b11814ce829c25e4", size = 44019, upload-time = "2026-07-20T13:43:00.335Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/8f/6f7273a7adb8d73fc8d21ede4376a3e475e52f98435c6007f69100dec8ca/bracex-3.0.1-py3-none-any.whl", hash = "sha256:6523ad83aeb5098a4ee597cff0f964442ff74e460bd3fafaffab6a013ff2288c", size = 11940, upload-time = "2026-07-20T13:42:59.268Z" }, +] + [[package]] name = "catalogue" version = "2.0.10" @@ -597,6 +606,7 @@ grpc = [ { name = "protobuf" }, ] server = [ + { name = "e2b" }, { name = "fastapi" }, { name = "graphon" }, { name = "jsonschema" }, @@ -627,6 +637,7 @@ docs = [ [package.metadata] requires-dist = [ + { name = "e2b", marker = "extra == 'server'", specifier = ">=2.34.0,<3.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, { name = "graphon", marker = "extra == 'server'", specifier = "==0.5.2" }, { name = "grpclib", extras = ["protobuf"], marker = "extra == 'grpc'", specifier = ">=0.4.9,<0.5.0" }, @@ -672,6 +683,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" }, ] +[[package]] +name = "dockerfile-parse" +version = "2.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/df/929ee0b5d2c8bd8d713c45e71b94ab57c7e11e322130724d54f469b2cd48/dockerfile-parse-2.0.1.tar.gz", hash = "sha256:3184ccdc513221983e503ac00e1aa504a2aa8f84e5de673c46b0b6eee99ec7bc", size = 24556, upload-time = "2023-07-18T13:36:07.897Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/6c/79cd5bc1b880d8c1a9a5550aa8dacd57353fa3bb2457227e1fb47383eb49/dockerfile_parse-2.0.1-py2.py3-none-any.whl", hash = "sha256:bdffd126d2eb26acf1066acb54cb2e336682e1d72b974a40894fac76a4df17f6", size = 14845, upload-time = "2023-07-18T13:36:06.052Z" }, +] + [[package]] name = "docstring-parser" version = "0.18.0" @@ -681,6 +701,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, ] +[[package]] +name = "e2b" +version = "2.34.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "dockerfile-parse" }, + { name = "h2" }, + { name = "httpcore" }, + { name = "httpx" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "python-dateutil" }, + { name = "rich" }, + { name = "typing-extensions" }, + { name = "wcmatch" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/22/5bcf34304b62e6984df9089eb996a4b2aa63618f901befe1f7cc73a24437/e2b-2.34.0.tar.gz", hash = "sha256:0dcc2694b3ead62e87b3d1070a05f66842dd2fa320fc1d7bd84c8fcfadeaf1b1", size = 189367, upload-time = "2026-07-17T09:59:23.844Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/ec/b2f9297afc25de7c27ab5a85db4060783680e2b132285e2498fdbad484cf/e2b-2.34.0-py3-none-any.whl", hash = "sha256:873323571d18bf633be45e59fc6271410b30dfbc81e8df85e711f4f184c03fea", size = 343447, upload-time = "2026-07-17T09:59:22.553Z" }, +] + [[package]] name = "emoji" version = "2.15.0" @@ -4085,6 +4127,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e3/bd/fa9bb053192491b3867ba07d2343d9f2252e00811567d30ae8d0f78136fe/watchfiles-1.1.1-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:a916a2932da8f8ab582f242c065f5c81bed3462849ca79ee357dd9551b0e9b01", size = 622112, upload-time = "2025-10-14T15:05:50.941Z" }, ] +[[package]] +name = "wcmatch" +version = "10.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "bracex" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/11/15/dc61746d8c0852f6d711ad09c774b63cf7c8211aa49e30871ac3d342b7e2/wcmatch-10.2.1.tar.gz", hash = "sha256:ecac70a5c70e62ba854b78318d3a1408e8651f8f1c96e5837743b71aa6a4fb92", size = 132497, upload-time = "2026-07-02T17:21:48.484Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/ba/20b48eedeab5316bf9a502bb9eb7e3b1588bd61d0f565822fefa8f06e10b/wcmatch-10.2.1-py3-none-any.whl", hash = "sha256:2d775395b93f233af66690f62cb9d52b084ec159a31cc4084f4069d72f437acd", size = 39763, upload-time = "2026-07-02T17:21:47.134Z" }, +] + [[package]] name = "weasel" version = "1.0.0" diff --git a/docker/.env.example b/docker/.env.example index 20ed14db271..c79c3b32334 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -222,6 +222,8 @@ PLUGIN_DAEMON_PORT=5002 PLUGIN_DAEMON_KEY=lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi PLUGIN_DAEMON_URL=http://plugin_daemon:5002 PLUGIN_MAX_PACKAGE_SIZE=52428800 +# Compose maps this byte value to DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES. +PLUGIN_MAX_FILE_SIZE=52428800 PLUGIN_MODEL_SCHEMA_CACHE_TTL=3600 PLUGIN_PPROF_ENABLED=false PLUGIN_DEBUGGING_HOST=0.0.0.0 @@ -269,8 +271,16 @@ DIFY_AGENT_PLUGIN_DAEMON_API_KEY= # DIFY_AGENT_INNER_API_KEY must match API/worker INNER_API_KEY_FOR_PLUGIN, not INNER_API_KEY. DIFY_AGENT_INNER_API_URL= DIFY_AGENT_INNER_API_KEY= -DIFY_AGENT_SHELLCTL_ENTRYPOINT=http://local_sandbox:5004 -DIFY_AGENT_SHELLCTL_AUTH_TOKEN= +# Select exactly one coherent Home Snapshot + Sandbox backend. +DIFY_AGENT_RUNTIME_BACKEND=local +DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://local_sandbox:5004 +DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN= +# E2B_API_KEY and E2B_API_TOKEN remain accepted as deployment-level fallbacks. +DIFY_AGENT_E2B_API_KEY= +DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox +DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600 +DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN= +DIFY_AGENT_E2B_SHELLCTL_PORT=5004 DIFY_AGENT_STUB_API_BASE_URL=http://agent_backend:5050/agent-stub # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index 9aebfd80f6e..89f95687e51 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -548,7 +548,7 @@ services: - path: ./envs/core-services/local-sandbox.env required: false environment: - - SHELLCTL_AUTH_TOKEN=${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-} + - SHELLCTL_AUTH_TOKEN=${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} - HTTP_PROXY=http://agent_ssrf_proxy:3128 - HTTPS_PROXY=http://agent_ssrf_proxy:3128 - NO_PROXY=localhost,127.0.0.1 @@ -668,8 +668,15 @@ services: DIFY_AGENT_PLUGIN_DAEMON_API_KEY: ${DIFY_AGENT_PLUGIN_DAEMON_API_KEY:-${PLUGIN_DAEMON_KEY:-lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi}} DIFY_AGENT_INNER_API_URL: ${DIFY_AGENT_INNER_API_URL:-${PLUGIN_DIFY_INNER_API_URL:-http://api:5001}} DIFY_AGENT_INNER_API_KEY: ${DIFY_AGENT_INNER_API_KEY:-${PLUGIN_DIFY_INNER_API_KEY:-QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1}} - DIFY_AGENT_SHELLCTL_ENTRYPOINT: ${DIFY_AGENT_SHELLCTL_ENTRYPOINT:-http://local_sandbox:5004} - DIFY_AGENT_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-} + DIFY_AGENT_RUNTIME_BACKEND: ${DIFY_AGENT_RUNTIME_BACKEND:-local} + DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT: ${DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT:-${DIFY_AGENT_SHELLCTL_ENTRYPOINT:-http://local_sandbox:5004}} + DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN: ${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} + DIFY_AGENT_E2B_API_KEY: ${DIFY_AGENT_E2B_API_KEY:-${E2B_API_KEY:-${E2B_API_TOKEN:-}}} + DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox} + DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS: ${DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS:-3600} + DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN:-} + DIFY_AGENT_E2B_SHELLCTL_PORT: ${DIFY_AGENT_E2B_SHELLCTL_PORT:-5004} + DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES: ${PLUGIN_MAX_FILE_SIZE:-52428800} DIFY_AGENT_STUB_API_BASE_URL: ${DIFY_AGENT_STUB_API_BASE_URL:-http://agent_backend:5050/agent-stub} # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. @@ -683,8 +690,6 @@ services: condition: service_started plugin_daemon: condition: service_started - local_sandbox: - condition: service_started networks: - default # Shared internal network with local_sandbox so agent_backend can reach it diff --git a/docker/docker-compose.e2b.yaml b/docker/docker-compose.e2b.yaml new file mode 100644 index 00000000000..2bfb5d566b5 --- /dev/null +++ b/docker/docker-compose.e2b.yaml @@ -0,0 +1,50 @@ +x-api-build: &api-build + context: .. + dockerfile: api/Dockerfile + +services: + api: + image: dify-api:e2b-local + build: *api-build + + api_websocket: + image: dify-api:e2b-local + build: *api-build + + worker: + image: dify-api:e2b-local + build: *api-build + + worker_beat: + image: dify-api:e2b-local + build: *api-build + + agent_backend: + image: dify-agent-backend:e2b-local + build: + context: .. + dockerfile: dify-agent/Dockerfile + environment: + DIFY_AGENT_RUNTIME_BACKEND: e2b + DIFY_AGENT_E2B_API_KEY: ${DIFY_AGENT_E2B_API_KEY:-${E2B_API_KEY:-${E2B_API_TOKEN:-}}} + DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox} + ports: + - "${EXPOSE_AGENT_BACKEND_PORT:-15050}:5050" + + plugin_daemon: + ports: !override [] + + # Keep the normal local sandbox out of an E2B deployment without changing + # the default compose stack. + local_sandbox: + profiles: + - local-agent-sandbox + + # This overlay is intended for branch validation and deliberately starts + # from an isolated PostgreSQL data volume. + db_postgres: + volumes: !override + - dify_e2b_postgres_data:/var/lib/postgresql/data + +volumes: + dify_e2b_postgres_data: diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 3e1384cc5ee..15f7ba73547 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -554,7 +554,7 @@ services: - path: ./envs/core-services/local-sandbox.env required: false environment: - - SHELLCTL_AUTH_TOKEN=${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-} + - SHELLCTL_AUTH_TOKEN=${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} - HTTP_PROXY=http://agent_ssrf_proxy:3128 - HTTPS_PROXY=http://agent_ssrf_proxy:3128 - NO_PROXY=localhost,127.0.0.1 @@ -674,8 +674,15 @@ services: DIFY_AGENT_PLUGIN_DAEMON_API_KEY: ${DIFY_AGENT_PLUGIN_DAEMON_API_KEY:-${PLUGIN_DAEMON_KEY:-lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi}} DIFY_AGENT_INNER_API_URL: ${DIFY_AGENT_INNER_API_URL:-${PLUGIN_DIFY_INNER_API_URL:-http://api:5001}} DIFY_AGENT_INNER_API_KEY: ${DIFY_AGENT_INNER_API_KEY:-${PLUGIN_DIFY_INNER_API_KEY:-QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1}} - DIFY_AGENT_SHELLCTL_ENTRYPOINT: ${DIFY_AGENT_SHELLCTL_ENTRYPOINT:-http://local_sandbox:5004} - DIFY_AGENT_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-} + DIFY_AGENT_RUNTIME_BACKEND: ${DIFY_AGENT_RUNTIME_BACKEND:-local} + DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT: ${DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT:-${DIFY_AGENT_SHELLCTL_ENTRYPOINT:-http://local_sandbox:5004}} + DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN: ${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} + DIFY_AGENT_E2B_API_KEY: ${DIFY_AGENT_E2B_API_KEY:-${E2B_API_KEY:-${E2B_API_TOKEN:-}}} + DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox} + DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS: ${DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS:-3600} + DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN:-} + DIFY_AGENT_E2B_SHELLCTL_PORT: ${DIFY_AGENT_E2B_SHELLCTL_PORT:-5004} + DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES: ${PLUGIN_MAX_FILE_SIZE:-52428800} DIFY_AGENT_STUB_API_BASE_URL: ${DIFY_AGENT_STUB_API_BASE_URL:-http://agent_backend:5050/agent-stub} # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. @@ -689,8 +696,6 @@ services: condition: service_started plugin_daemon: condition: service_started - local_sandbox: - condition: service_started networks: - default # Shared internal network with local_sandbox so agent_backend can reach it diff --git a/docker/envs/core-services/dify-agent.env.example b/docker/envs/core-services/dify-agent.env.example index f2d7155f521..d434ecf5945 100644 --- a/docker/envs/core-services/dify-agent.env.example +++ b/docker/envs/core-services/dify-agent.env.example @@ -22,8 +22,19 @@ DIFY_AGENT_PLUGIN_DAEMON_API_KEY= DIFY_AGENT_INNER_API_URL= DIFY_AGENT_INNER_API_KEY= -DIFY_AGENT_SHELLCTL_ENTRYPOINT=http://local_sandbox:5004 -DIFY_AGENT_SHELLCTL_AUTH_TOKEN= +# Select exactly one coherent Home Snapshot + Sandbox backend. +DIFY_AGENT_RUNTIME_BACKEND=local +DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://local_sandbox:5004 +DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN= +# E2B_API_KEY and E2B_API_TOKEN remain accepted as deployment-level fallbacks. +DIFY_AGENT_E2B_API_KEY= +DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox +DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600 +DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN= +DIFY_AGENT_E2B_SHELLCTL_PORT=5004 +# Standalone/direct-container byte limit. Full Docker Compose derives this from +# PLUGIN_MAX_FILE_SIZE in docker/.env. +DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES=52428800 DIFY_AGENT_STUB_API_BASE_URL=http://agent_backend:5050/agent-stub # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. diff --git a/packages/contracts/generated/api/console/agent/orpc.gen.ts b/packages/contracts/generated/api/console/agent/orpc.gen.ts index 12b58ef4188..2cc590c433d 100644 --- a/packages/contracts/generated/api/console/agent/orpc.gen.ts +++ b/packages/contracts/generated/api/console/agent/orpc.gen.ts @@ -143,7 +143,6 @@ import { zPostAgentByAgentIdCopyBody, zPostAgentByAgentIdCopyPath, zPostAgentByAgentIdCopyResponse, - zPostAgentByAgentIdDebugConversationRefreshBody, zPostAgentByAgentIdDebugConversationRefreshPath, zPostAgentByAgentIdDebugConversationRefreshResponse, zPostAgentByAgentIdFeaturesBody, @@ -853,12 +852,7 @@ export const post12 = oc path: '/agent/{agent_id}/debug-conversation/refresh', tags: ['console'], }) - .input( - z.object({ - body: zPostAgentByAgentIdDebugConversationRefreshBody.optional(), - params: zPostAgentByAgentIdDebugConversationRefreshPath, - }), - ) + .input(z.object({ params: zPostAgentByAgentIdDebugConversationRefreshPath })) .output(zPostAgentByAgentIdDebugConversationRefreshResponse) export const refresh = { diff --git a/packages/contracts/generated/api/console/agent/types.gen.ts b/packages/contracts/generated/api/console/agent/types.gen.ts index 0a2bae65ac1..f854fe67ac7 100644 --- a/packages/contracts/generated/api/console/agent/types.gen.ts +++ b/packages/contracts/generated/api/console/agent/types.gen.ts @@ -282,10 +282,6 @@ export type AgentAppCopyPayload = { role?: string | null } -export type AgentDebugConversationRefreshPayload = { - draft_type?: AgentConfigDraftType -} - export type AgentDebugConversationRefreshResponse = { debug_conversation_has_messages?: boolean debug_conversation_id: string @@ -426,7 +422,6 @@ export type AgentReferencingWorkflowsResponse = { } export type SandboxInfoResponse = { - session_id: string workspace_cwd: string } @@ -445,7 +440,8 @@ export type SandboxReadResponse = { } export type AgentSandboxUploadPayload = { - conversation_id: string + caller_id: string + caller_type: 'build_draft' | 'conversation' path: string } @@ -822,8 +818,6 @@ export type AgentConfigSkillMarkdownResponse = { truncated: boolean } -export type AgentConfigDraftType = 'debug_build' | 'draft' - export type AgentDriveItemResponse = { created_at?: number | null file_kind: string @@ -1227,6 +1221,8 @@ export type AgentSoulToolsConfig = { dify_tools?: Array } +export type AgentConfigDraftType = 'debug_build' | 'draft' + export type DeclaredOutputConfig = { array_item?: DeclaredArrayItem | null check?: DeclaredOutputCheckConfig | null @@ -2756,7 +2752,7 @@ export type PostAgentByAgentIdCopyResponse = PostAgentByAgentIdCopyResponses[keyof PostAgentByAgentIdCopyResponses] export type PostAgentByAgentIdDebugConversationRefreshData = { - body?: AgentDebugConversationRefreshPayload + body?: never path: { agent_id: string } @@ -3077,7 +3073,8 @@ export type GetAgentByAgentIdSandboxData = { agent_id: string } query: { - conversation_id: string + caller_id: string + caller_type: 'build_draft' | 'conversation' } url: '/agent/{agent_id}/sandbox' } @@ -3095,7 +3092,8 @@ export type GetAgentByAgentIdSandboxFilesData = { agent_id: string } query: { - conversation_id: string + caller_id: string + caller_type: 'build_draft' | 'conversation' path?: string } url: '/agent/{agent_id}/sandbox/files' @@ -3114,7 +3112,8 @@ export type GetAgentByAgentIdSandboxFilesReadData = { agent_id: string } query: { - conversation_id: string + caller_id: string + caller_type: 'build_draft' | 'conversation' path: string } url: '/agent/{agent_id}/sandbox/files/read' diff --git a/packages/contracts/generated/api/console/agent/zod.gen.ts b/packages/contracts/generated/api/console/agent/zod.gen.ts index dfa6917bfb7..08c7c93377a 100644 --- a/packages/contracts/generated/api/console/agent/zod.gen.ts +++ b/packages/contracts/generated/api/console/agent/zod.gen.ts @@ -196,7 +196,6 @@ export const zAgentPublishPayload = z.object({ * SandboxInfoResponse */ export const zSandboxInfoResponse = z.object({ - session_id: z.string(), workspace_cwd: z.string(), }) @@ -215,7 +214,8 @@ export const zSandboxReadResponse = z.object({ * AgentSandboxUploadPayload */ export const zAgentSandboxUploadPayload = z.object({ - conversation_id: z.string().min(1), + caller_id: z.string().min(1), + caller_type: z.enum(['build_draft', 'conversation']), path: z.string().min(1), }) @@ -650,45 +650,6 @@ export const zAgentConfigSkillInspectResponse = z.object({ warnings: z.array(z.string()).optional(), }) -/** - * AgentConfigDraftType - * - * Editable Agent Soul draft workspace type. - */ -export const zAgentConfigDraftType = z.enum(['debug_build', 'draft']) - -/** - * AgentDebugConversationRefreshPayload - */ -export const zAgentDebugConversationRefreshPayload = z.object({ - draft_type: zAgentConfigDraftType.optional().default('debug_build'), -}) - -/** - * AgentConfigDraftSummaryResponse - */ -export const zAgentConfigDraftSummaryResponse = z.object({ - account_id: z.string().nullish(), - agent_id: z.string(), - base_snapshot_id: z.string().nullish(), - created_at: z.int().nullish(), - created_by: z.string().nullish(), - draft_type: zAgentConfigDraftType, - id: z.string(), - updated_at: z.int().nullish(), - updated_by: z.string().nullish(), -}) - -/** - * AgentPublishResponse - */ -export const zAgentPublishResponse = z.object({ - active_config_snapshot: zAgentConfigSnapshotSummaryResponse.nullish(), - active_config_snapshot_id: z.string(), - draft: zAgentConfigDraftSummaryResponse.nullish(), - result: z.string(), -}) - /** * AgentDriveItemResponse */ @@ -1278,6 +1239,38 @@ export const zAgentSoulPromptConfig = z.object({ system_prompt: z.string().optional().default(''), }) +/** + * AgentConfigDraftType + * + * Editable Agent Soul draft workspace type. + */ +export const zAgentConfigDraftType = z.enum(['debug_build', 'draft']) + +/** + * AgentConfigDraftSummaryResponse + */ +export const zAgentConfigDraftSummaryResponse = z.object({ + account_id: z.string().nullish(), + agent_id: z.string(), + base_snapshot_id: z.string().nullish(), + created_at: z.int().nullish(), + created_by: z.string().nullish(), + draft_type: zAgentConfigDraftType, + id: z.string(), + updated_at: z.int().nullish(), + updated_by: z.string().nullish(), +}) + +/** + * AgentPublishResponse + */ +export const zAgentPublishResponse = z.object({ + active_config_snapshot: zAgentConfigSnapshotSummaryResponse.nullish(), + active_config_snapshot_id: z.string(), + draft: zAgentConfigDraftSummaryResponse.nullish(), + result: z.string(), +}) + /** * AgentHumanContactConfig */ @@ -3263,8 +3256,6 @@ export const zPostAgentByAgentIdCopyPath = z.object({ */ export const zPostAgentByAgentIdCopyResponse = zAgentAppDetailWithSite -export const zPostAgentByAgentIdDebugConversationRefreshBody = zAgentDebugConversationRefreshPayload - export const zPostAgentByAgentIdDebugConversationRefreshPath = z.object({ agent_id: z.uuid(), }) @@ -3471,7 +3462,8 @@ export const zGetAgentByAgentIdSandboxPath = z.object({ }) export const zGetAgentByAgentIdSandboxQuery = z.object({ - conversation_id: z.string().min(1), + caller_id: z.string().min(1), + caller_type: z.enum(['build_draft', 'conversation']), }) /** @@ -3484,7 +3476,8 @@ export const zGetAgentByAgentIdSandboxFilesPath = z.object({ }) export const zGetAgentByAgentIdSandboxFilesQuery = z.object({ - conversation_id: z.string().min(1), + caller_id: z.string().min(1), + caller_type: z.enum(['build_draft', 'conversation']), path: z.string().optional().default('.'), }) @@ -3498,7 +3491,8 @@ export const zGetAgentByAgentIdSandboxFilesReadPath = z.object({ }) export const zGetAgentByAgentIdSandboxFilesReadQuery = z.object({ - conversation_id: z.string().min(1), + caller_id: z.string().min(1), + caller_type: z.enum(['build_draft', 'conversation']), path: z.string().min(1), }) diff --git a/packages/contracts/generated/api/console/apps/orpc.gen.ts b/packages/contracts/generated/api/console/apps/orpc.gen.ts index b84db3cbb4e..9f0a504d9be 100644 --- a/packages/contracts/generated/api/console/apps/orpc.gen.ts +++ b/packages/contracts/generated/api/console/apps/orpc.gen.ts @@ -3093,8 +3093,7 @@ export const get59 = oc .input( z.object({ params: zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFilesPath, - query: - zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFilesQuery.optional(), + query: zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFilesQuery, }), ) .output(zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFilesResponse) diff --git a/packages/contracts/generated/api/console/apps/types.gen.ts b/packages/contracts/generated/api/console/apps/types.gen.ts index c8bd2dfe64e..cfff4e2be85 100644 --- a/packages/contracts/generated/api/console/apps/types.gen.ts +++ b/packages/contracts/generated/api/console/apps/types.gen.ts @@ -874,7 +874,7 @@ export type SandboxReadResponse = { } export type WorkflowAgentSandboxUploadPayload = { - node_execution_id?: string | null + node_execution_id: string path: string } @@ -5658,8 +5658,8 @@ export type GetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFi node_id: string workflow_run_id: string } - query?: { - node_execution_id?: string + query: { + node_execution_id: string path?: string } url: '/apps/{app_id}/workflow-runs/{workflow_run_id}/agent-nodes/{node_id}/sandbox/files' @@ -5680,7 +5680,7 @@ export type GetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFi workflow_run_id: string } query: { - node_execution_id?: string + node_execution_id: string path: string } url: '/apps/{app_id}/workflow-runs/{workflow_run_id}/agent-nodes/{node_id}/sandbox/files/read' diff --git a/packages/contracts/generated/api/console/apps/zod.gen.ts b/packages/contracts/generated/api/console/apps/zod.gen.ts index 534a11ff431..8bdbda8172a 100644 --- a/packages/contracts/generated/api/console/apps/zod.gen.ts +++ b/packages/contracts/generated/api/console/apps/zod.gen.ts @@ -560,7 +560,7 @@ export const zSandboxReadResponse = z.object({ * WorkflowAgentSandboxUploadPayload */ export const zWorkflowAgentSandboxUploadPayload = z.object({ - node_execution_id: z.string().nullish(), + node_execution_id: z.string().min(1), path: z.string().min(1), }) @@ -5931,7 +5931,7 @@ export const zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandbox export const zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFilesQuery = z.object({ - node_execution_id: z.string().optional(), + node_execution_id: z.string().min(1), path: z.string().optional().default('.'), }) @@ -5950,7 +5950,7 @@ export const zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandbox export const zGetAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFilesReadQuery = z.object({ - node_execution_id: z.string().optional(), + node_execution_id: z.string().min(1), path: z.string().min(1), }) diff --git a/web/app/components/workflow/nodes/agent-v2/components/__tests__/agent-orchestrate-panel-content.spec.tsx b/web/app/components/workflow/nodes/agent-v2/components/__tests__/agent-orchestrate-panel-content.spec.tsx index ea410bd8635..714454e5298 100644 --- a/web/app/components/workflow/nodes/agent-v2/components/__tests__/agent-orchestrate-panel-content.spec.tsx +++ b/web/app/components/workflow/nodes/agent-v2/components/__tests__/agent-orchestrate-panel-content.spec.tsx @@ -302,6 +302,7 @@ vi.mock('@/service/client', async () => { }, files: { get: { + key: () => ['agent-sandbox-files'], queryOptions: () => ({ queryKey: ['sandbox-files'], queryFn: () => @@ -471,7 +472,9 @@ describe('WorkflowInlineAgentConfigureWorkspace', () => { mocks.loadBuildDraft.mockRejectedValue(new Response(null, { status: 404 })) mocks.checkoutBuildDraft.mockResolvedValue({ agent_soul: {}, - draft: {}, + draft: { + id: 'build-draft-1', + }, variant: 'agent_app', }) mocks.deleteBuildDraft.mockResolvedValue({ result: 'success' }) @@ -493,7 +496,9 @@ describe('WorkflowInlineAgentConfigureWorkspace', () => { system_prompt: 'Help with workflow tasks.', }, }, - draft: {}, + draft: { + id: 'build-draft-1', + }, variant: 'agent_app', }) mocks.saveDraft.mockResolvedValue(createInlineComposerState()) @@ -856,6 +861,12 @@ describe('WorkflowInlineAgentConfigureWorkspace', () => { }), }) + expect( + screen.queryByRole('button', { + name: 'agentV2.agentDetail.configure.workingDirectory.open', + }), + ).not.toBeInTheDocument() + fireEvent.click( await screen.findByRole('button', { name: 'send build message', @@ -877,11 +888,6 @@ describe('WorkflowInlineAgentConfigureWorkspace', () => { name: 'complete build conversation', }), ) - fireEvent.click( - await screen.findByRole('button', { - name: 'send build message', - }), - ) expect( screen.getByRole('button', { name: 'agentV2.agentDetail.configure.workingDirectory.open', diff --git a/web/app/components/workflow/nodes/agent-v2/components/agent-orchestrate-panel-content.tsx b/web/app/components/workflow/nodes/agent-v2/components/agent-orchestrate-panel-content.tsx index 0c9c753be86..e07c91a8fd6 100644 --- a/web/app/components/workflow/nodes/agent-v2/components/agent-orchestrate-panel-content.tsx +++ b/web/app/components/workflow/nodes/agent-v2/components/agent-orchestrate-panel-content.tsx @@ -303,18 +303,18 @@ function WorkflowInlineAgentConfigureWorkspaceContent({ const [completedBuildConversationId, setCompletedBuildConversationId] = useState( null, ) - const [workflowRunId, setWorkflowRunId] = useState(null) const rightPanelChatControllerRef = useRef(null) const appId = flowType === FlowType.appFlow ? flowId : undefined const conversationIds = useAtomValue(agentConfigureConversationIdsAtom) const [rightPanelMode, setRightPanelMode] = useAtom(agentConfigureRightPanelModeAtom) const previewEnabled = systemFeatures?.deployment_edition !== 'COMMUNITY' const workingDirectoryPanel = useAgentWorkingDirectoryPanel({ + type: 'agent', agentId, - appId, - conversationId: conversationIds[rightPanelMode], - nodeId, - workflowRunId, + caller: { + type: 'build_draft', + id: buildDraft.id, + }, }) const resetConversation = useSetAtom(resetAgentConfigureConversationAtom) const setConversationId = useSetAtom(setAgentConfigureConversationIdAtom) @@ -446,9 +446,8 @@ function WorkflowInlineAgentConfigureWorkspaceContent({ await refreshDebugConversationAsync().catch(() => undefined) setCompletedBuildConversationId(null) setConversationId({ mode: 'build', conversationId: null }) - setWorkflowRunId(null) setClearPreviewChat(true) - }, [refreshDebugConversationAsync, setClearPreviewChat, setConversationId, setWorkflowRunId]) + }, [refreshDebugConversationAsync, setClearPreviewChat, setConversationId]) const rebaseComposerDraftFromSoulConfig = useCallback( (agentSoulConfig?: AgentSoulConfig) => { rebaseComposerDraft({ @@ -655,7 +654,8 @@ function WorkflowInlineAgentConfigureWorkspaceContent({ (conversationIds.build === completedBuildConversationId || (conversationIds.build === inlineComposerState?.debug_conversation_id && (inlineComposerState?.debug_conversation_has_messages ?? false))) - const showWorkingDirectoryAction = rightPanelMode === 'build' && buildConversationHasAgentResponse + const showWorkingDirectoryAction = + rightPanelMode === 'build' && !!buildDraft.id && buildConversationHasAgentResponse const restartCurrentChat = () => { if (isRestartCurrentChatDisabled) return @@ -760,18 +760,13 @@ function WorkflowInlineAgentConfigureWorkspaceContent({ conversationIds={conversationIds} mode={rightPanelMode} onClearChatListChange={setClearPreviewChat} - onConversationComplete={(mode, completedConversationId, completedWorkflowRunId) => { + onConversationComplete={(mode, completedConversationId) => { if (mode !== 'build' || !isBuildCallbackCurrent(buildCallbackGeneration)) return setCompletedBuildConversationId(completedConversationId) - setWorkflowRunId(completedWorkflowRunId ?? completedConversationId) invalidateAgentWorkingDirectoryFiles({ - agentId, - appId, conversationId: completedConversationId, - nodeId, queryClient, - workflowRunId: completedWorkflowRunId ?? completedConversationId, }) buildDraftActions.refreshBuildDraftAfterBuildChat(() => finishBuildAction(buildCallbackGeneration), @@ -781,14 +776,9 @@ function WorkflowInlineAgentConfigureWorkspaceContent({ if (mode === 'build' && !isBuildCallbackCurrent(buildCallbackGeneration)) return setConversationId({ mode, conversationId }) }} - onWorkflowRunIdChange={(nextWorkflowRunId) => { - if (!isBuildCallbackCurrent(buildCallbackGeneration)) return - if (nextWorkflowRunId) setWorkflowRunId(nextWorkflowRunId) - }} onSaveDraftBeforeRun={ rightPanelMode === 'build' ? () => { - setWorkflowRunId(null) return runBuildPreparation({ generation: buildCallbackGeneration, markBuildChatStarted: true, diff --git a/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx b/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx index 08152cbd89e..b91887b5797 100644 --- a/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx +++ b/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx @@ -2040,7 +2040,9 @@ describe('AgentConfigurePage', () => { system_prompt: 'build prompt', }, }, - draft: {}, + draft: { + id: 'build-draft-1', + }, variant: 'agent_app', }, dataUpdatedAt: 1, diff --git a/web/features/agent-v2/agent-detail/configure/components/composer-session.tsx b/web/features/agent-v2/agent-detail/configure/components/composer-session.tsx index 15964340ff5..0086f999730 100644 --- a/web/features/agent-v2/agent-detail/configure/components/composer-session.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/composer-session.tsx @@ -297,10 +297,25 @@ function AgentConfigurePageComposerContent({ const rightPanelChatControllerRef = useRef(null) const conversationIds = useAtomValue(agentConfigureConversationIdsAtom) const rightPanelChatMode = rightPanelMode - const workingDirectoryPanel = useAgentWorkingDirectoryPanel({ - agentId, - conversationId: conversationIds[rightPanelChatMode], - }) + const workingDirectoryPanel = useAgentWorkingDirectoryPanel( + rightPanelChatMode === 'build' + ? { + type: 'agent', + agentId, + caller: { + type: 'build_draft', + id: buildDraft.id, + }, + } + : { + type: 'agent', + agentId, + caller: { + type: 'conversation', + id: conversationIds.preview, + }, + }, + ) const showChatFeatures = useAtomValue(agentConfigureShowChatFeaturesAtom) const showPreviewVersions = useAtomValue(agentConfigureShowPreviewVersionsAtom) const resetConversation = useSetAtom(resetAgentConfigureConversationAtom) @@ -422,7 +437,7 @@ function AgentConfigurePageComposerContent({ (conversationIds.build === agentQuery.data?.debug_conversation_id && (agentQuery.data?.debug_conversation_has_messages ?? false))) const showWorkingDirectoryAction = - rightPanelChatMode === 'build' && buildConversationHasAgentResponse + rightPanelChatMode === 'build' && !!buildDraft.id && buildConversationHasAgentResponse const restartCurrentChat = () => { if (isRestartCurrentChatDisabled) return @@ -543,7 +558,6 @@ function AgentConfigurePageComposerContent({ setCompletedBuildConversationId(completedConversationId) invalidateAgentWorkingDirectoryFiles({ - agentId, conversationId: completedConversationId, queryClient, }) diff --git a/web/features/agent-v2/agent-detail/configure/components/preview/__tests__/working-directory-panel.spec.tsx b/web/features/agent-v2/agent-detail/configure/components/preview/__tests__/working-directory-panel.spec.tsx index bda6f3ae5b5..4f5dcfb924d 100644 --- a/web/features/agent-v2/agent-detail/configure/components/preview/__tests__/working-directory-panel.spec.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/preview/__tests__/working-directory-panel.spec.tsx @@ -160,13 +160,38 @@ function renderWorkingDirectoryPanel() { source={{ type: 'agent', agentId: 'agent-1', - conversationId: 'conversation-1', + callerType: 'conversation', + callerId: 'conversation-1', }} /> , ) } +function renderWorkflowWorkingDirectoryPanel() { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, + }) + const rendered = render( + + + , + ) + return { ...rendered, queryClient } +} + describe('AgentWorkingDirectoryPanel', () => { beforeEach(() => { vi.clearAllMocks() @@ -197,6 +222,24 @@ describe('AgentWorkingDirectoryPanel', () => { truncated: false, }), })) + mocks.workflowSandboxFilesQueryOptions.mockImplementation(({ input }: QueryOptionsInput) => ({ + queryKey: ['workflow-sandbox-files', input], + queryFn: async () => ({ + path: input.query?.path ?? '.', + entries: [{ name: 'chart.png', type: 'file' }], + }), + })) + mocks.workflowSandboxFileReadQueryOptions.mockImplementation( + ({ input }: QueryOptionsInput) => ({ + queryKey: ['workflow-sandbox-file-read', input], + queryFn: async () => ({ + binary: false, + path: input.query?.path ?? '', + text: null, + truncated: false, + }), + }), + ) }) it('should download the selected working directory file from the preview header download action', async () => { @@ -226,7 +269,8 @@ describe('AgentWorkingDirectoryPanel', () => { agent_id: 'agent-1', }, body: { - conversation_id: 'conversation-1', + caller_type: 'conversation', + caller_id: 'conversation-1', path: '~/workspace/notes.md', }, }) @@ -269,7 +313,8 @@ describe('AgentWorkingDirectoryPanel', () => { agent_id: 'agent-1', }, body: { - conversation_id: 'conversation-1', + caller_type: 'conversation', + caller_id: 'conversation-1', path: '~/workspace/model.bin', }, }) @@ -301,7 +346,8 @@ describe('AgentWorkingDirectoryPanel', () => { agent_id: 'agent-1', }, body: { - conversation_id: 'conversation-1', + caller_type: 'conversation', + caller_id: 'conversation-1', path: '~/workspace/chart.png', }, }) @@ -310,4 +356,28 @@ describe('AgentWorkingDirectoryPanel', () => { ).not.toBeInTheDocument() expect(mocks.downloadUrl).not.toHaveBeenCalled() }) + + it('should scope workflow image previews to the exact node execution', async () => { + const { queryClient } = renderWorkflowWorkingDirectoryPanel() + + await waitFor(() => { + expect(mocks.workflowSandboxFileUploadClientPost).toHaveBeenCalled() + }) + + expect( + queryClient.getQueryCache().find({ + queryKey: [ + 'agent-v2', + 'working-directory', + 'image-preview', + 'workflow-node', + 'app-1', + 'run-1', + 'node-1', + 'execution-1', + 'chart.png', + ], + }), + ).toBeDefined() + }) }) diff --git a/web/features/agent-v2/agent-detail/configure/components/preview/chat-history.ts b/web/features/agent-v2/agent-detail/configure/components/preview/chat-history.ts index 66e5d7532f0..6f3694dbd0f 100644 --- a/web/features/agent-v2/agent-detail/configure/components/preview/chat-history.ts +++ b/web/features/agent-v2/agent-detail/configure/components/preview/chat-history.ts @@ -97,15 +97,6 @@ const toFeedback = ( } } -export function getLastWorkflowRunId(messages: MessageDetailResponse[]) { - for (let index = messages.length - 1; index >= 0; index--) { - const workflowRunId = messages[index]?.workflow_run_id - if (workflowRunId) return workflowRunId - } - - return null -} - export function getFormattedAgentDebugChatTree( messages: MessageDetailResponse[], ): ChatItemInTree[] { diff --git a/web/features/agent-v2/agent-detail/configure/components/preview/chat-runtime.tsx b/web/features/agent-v2/agent-detail/configure/components/preview/chat-runtime.tsx index 0da792b4c20..ba38040d440 100644 --- a/web/features/agent-v2/agent-detail/configure/components/preview/chat-runtime.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/preview/chat-runtime.tsx @@ -5,10 +5,10 @@ import type { ReactNode, Ref } from 'react' import type { AgentChatMessageSender, AgentPreviewChatController } from './chat-conversation' import type { AnswerActionPosition } from '@/app/components/base/chat/chat/answer/operation' import { skipToken, useQuery } from '@tanstack/react-query' -import { useCallback, useEffect, useMemo, useState } from 'react' +import { useCallback, useMemo, useState } from 'react' import Loading from '@/app/components/base/loading' import { consoleQuery } from '@/service/client' -import { getFormattedAgentDebugChatTree, getLastWorkflowRunId } from './chat-history' +import { getFormattedAgentDebugChatTree } from './chat-history' import { AgentPreviewChatSession } from './chat-session' export type AgentChatRuntimeEmptyStateProps = { @@ -41,7 +41,6 @@ export type AgentChatRuntimeProps = { onClearChatListChange: (clearChatList: boolean) => void onConversationComplete?: (conversationId: string, workflowRunId?: string) => void onConversationIdChange?: (conversationId: string) => void - onWorkflowRunIdChange?: (workflowRunId: string | null) => void onBeforeSpeechToText?: () => Promise onSaveDraftBeforeRun?: () => Promise onSendInterrupted?: () => void @@ -69,7 +68,6 @@ export function AgentChatRuntime({ onClearChatListChange, onConversationComplete, onConversationIdChange, - onWorkflowRunIdChange, onBeforeSpeechToText, onSendInterrupted, onSaveDraftBeforeRun, @@ -104,12 +102,6 @@ export function AgentChatRuntime({ () => getFormattedAgentDebugChatTree(historyQuery.data?.data ?? []), [historyQuery.data?.data], ) - useEffect(() => { - if (!conversationId || !historyQuery.data) return - - onWorkflowRunIdChange?.(getLastWorkflowRunId(historyQuery.data.data ?? [])) - }, [conversationId, historyQuery.data, onWorkflowRunIdChange]) - if (conversationId && historyQuery.isPending && !conversationBelongsToCurrentSession) { return (
diff --git a/web/features/agent-v2/agent-detail/configure/components/preview/hook/use-working-directory-panel.tsx b/web/features/agent-v2/agent-detail/configure/components/preview/hook/use-working-directory-panel.tsx index b0ec68ddfe2..dbdd3bf43cf 100644 --- a/web/features/agent-v2/agent-detail/configure/components/preview/hook/use-working-directory-panel.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/preview/hook/use-working-directory-panel.tsx @@ -13,12 +13,10 @@ export function invalidateAgentWorkingDirectoryFiles({ nodeId, queryClient, }: { - agentId: string appId?: string conversationId?: string | null nodeId?: string queryClient: QueryClient - workflowRunId?: string | null }) { if (appId && nodeId) { void queryClient.invalidateQueries({ @@ -37,44 +35,65 @@ export function invalidateAgentWorkingDirectoryFiles({ }) } -export function useAgentWorkingDirectoryPanel({ - agentId, - appId, - conversationId, - nodeId, - workflowRunId, -}: { - agentId: string - appId?: string - conversationId?: string | null - nodeId?: string - workflowRunId?: string | null -}): { +type AgentWorkingDirectoryPanelInput = + | { + type: 'agent' + agentId: string + caller: + | { + type: 'build_draft' + id?: string + } + | { + type: 'conversation' + id?: string | null + } + } + | { + type: 'workflow-node' + appId?: string + nodeId: string + nodeExecutionId?: string + workflowRunId?: string | null + } + +export function useAgentWorkingDirectoryPanel(input: AgentWorkingDirectoryPanelInput): { closeWorkingDirectory: () => void openWorkingDirectory: () => void panel: ReactNode } { const [open, setOpen] = useState(false) - const source: AgentWorkingDirectorySource = - appId && nodeId + let source: AgentWorkingDirectorySource | undefined + if (input.type === 'workflow-node') { + const { appId, nodeExecutionId, nodeId, workflowRunId } = input + source = + appId && nodeExecutionId && workflowRunId + ? { + type: 'workflow-node', + appId, + nodeId, + nodeExecutionId, + workflowRunId, + } + : undefined + } else { + const callerId = input.caller.id + source = callerId ? { - type: 'workflow-node', - appId, - conversationId, - nodeId, - workflowRunId, - } - : { type: 'agent', - agentId, - conversationId, + agentId: input.agentId, + callerType: input.caller.type, + callerId, } + : undefined + } return { closeWorkingDirectory: () => setOpen(false), openWorkingDirectory: () => setOpen(true), - panel: open ? ( - - ) : null, + panel: + open && source ? ( + + ) : null, } } diff --git a/web/features/agent-v2/agent-detail/configure/components/preview/working-directory-panel.tsx b/web/features/agent-v2/agent-detail/configure/components/preview/working-directory-panel.tsx index ec5f6f9b61b..f7c37395ea2 100644 --- a/web/features/agent-v2/agent-detail/configure/components/preview/working-directory-panel.tsx +++ b/web/features/agent-v2/agent-detail/configure/components/preview/working-directory-panel.tsx @@ -30,14 +30,15 @@ export type AgentWorkingDirectorySource = | { type: 'agent' agentId: string - conversationId?: string | null + callerType: 'conversation' | 'build_draft' + callerId: string } | { type: 'workflow-node' - appId?: string - conversationId?: string | null + appId: string nodeId: string - workflowRunId?: string | null + nodeExecutionId: string + workflowRunId: string } type SandboxErrorPayload = { @@ -211,12 +212,12 @@ function countReadableFiles(files: AgentFileNode[]): number { }, 0) } -async function isNoActiveSessionError(error: unknown) { +async function isNoActiveBindingError(error: unknown) { if (!(error instanceof Response) || error.status !== 404) return false try { const payload = (await error.clone().json()) as SandboxErrorPayload - return payload.code === 'no_active_session' + return payload.code === 'no_active_binding' } catch { return false } @@ -249,19 +250,16 @@ export function AgentWorkingDirectoryPanel({ const [pendingOpenFolderPaths, setPendingOpenFolderPaths] = useState([]) const [downloadActionLoadingTarget, setDownloadActionLoadingTarget] = useState(null) - const workflowNodeRunId = - source.type === 'workflow-node' ? (source.workflowRunId ?? source.conversationId) : undefined - const hasWorkingDirectorySource = - source.type === 'agent' ? !!source.conversationId : !!source.appId && !!workflowNodeRunId const sandboxInfoQueryOptions = consoleQuery.agent.byAgentId.sandbox.get.queryOptions({ input: - source.type === 'agent' && source.conversationId + source.type === 'agent' ? { params: { agent_id: source.agentId, }, query: { - conversation_id: source.conversationId, + caller_type: source.callerType, + caller_id: source.callerId, }, } : skipToken, @@ -271,11 +269,10 @@ export function AgentWorkingDirectoryPanel({ }) const sandboxInfoQuery = useQuery({ ...sandboxInfoQueryOptions, - enabled: open && source.type === 'agent' && !!source.conversationId, + enabled: open && source.type === 'agent', retry: false, }) - const isSandboxInfoLoading = - source.type === 'agent' && !!source.conversationId && sandboxInfoQuery.isPending + const isSandboxInfoLoading = source.type === 'agent' && sandboxInfoQuery.isPending const workspaceDirectoryPath = sandboxInfoQuery.data?.workspace_cwd const directoryPath = selectedDirectoryPath ?? workspaceDirectoryPath ?? '.' const showReturnToWorkspaceButton = @@ -283,37 +280,35 @@ export function AgentWorkingDirectoryPanel({ const getFileListQueryOptions = (path: string) => source.type === 'agent' ? consoleQuery.agent.byAgentId.sandbox.files.get.queryOptions({ - input: - source.conversationId && !isSandboxInfoLoading - ? { - params: { - agent_id: source.agentId, - }, - query: { - conversation_id: source.conversationId, - path: toSandboxApiPath(path), - }, - } - : skipToken, + input: !isSandboxInfoLoading + ? { + params: { + agent_id: source.agentId, + }, + query: { + caller_type: source.callerType, + caller_id: source.callerId, + path: toSandboxApiPath(path), + }, + } + : skipToken, context: { silent: true, }, }) : consoleQuery.apps.byAppId.workflowRuns.byWorkflowRunId.agentNodes.byNodeId.sandbox.files.get.queryOptions( { - input: - source.appId && workflowNodeRunId - ? { - params: { - app_id: source.appId, - workflow_run_id: workflowNodeRunId, - node_id: source.nodeId, - }, - query: { - path: toSandboxApiPath(path), - }, - } - : skipToken, + input: { + params: { + app_id: source.appId, + workflow_run_id: source.workflowRunId, + node_id: source.nodeId, + }, + query: { + node_execution_id: source.nodeExecutionId, + path: toSandboxApiPath(path), + }, + }, context: { silent: true, }, @@ -333,7 +328,7 @@ export function AgentWorkingDirectoryPanel({ try { return await fileListQueryOptions.queryFn(context) } catch (error) { - if (await isNoActiveSessionError(error)) { + if (await isNoActiveBindingError(error)) { return { entries: [], path: '.', @@ -346,30 +341,28 @@ export function AgentWorkingDirectoryPanel({ retry: false, }) const expandedFolderQueries = useQueries({ - queries: hasWorkingDirectorySource - ? loadedFolderPaths.map((path) => { - const queryOptions = getFileListQueryOptions(path) + queries: loadedFolderPaths.map((path) => { + const queryOptions = getFileListQueryOptions(path) - return { - ...queryOptions, - queryFn: async (context): Promise => { - try { - return await queryOptions.queryFn(context) - } catch (error) { - if (await isNoActiveSessionError(error)) { - return { - entries: [], - path, - } - } - - throw error + return { + ...queryOptions, + queryFn: async (context): Promise => { + try { + return await queryOptions.queryFn(context) + } catch (error) { + if (await isNoActiveBindingError(error)) { + return { + entries: [], + path, } - }, - retry: false, + } + + throw error } - }) - : [], + }, + retry: false, + } + }), }) const workingDirectoryFiles = expandedFolderQueries.reduce( (files, query, index) => { @@ -386,8 +379,7 @@ export function AgentWorkingDirectoryPanel({ const selectedWorkingDirectoryFile = findReadableFile(workingDirectoryFiles, selectedFileId) ?? findFirstReadableFile(workingDirectoryFiles) - const isFileListLoading = - hasWorkingDirectorySource && (isSandboxInfoLoading || fileListQuery.isPending) + const isFileListLoading = isSandboxInfoLoading || fileListQuery.isPending const loadingFolderPaths = new Set( loadedFolderPaths.filter((path, index) => expandedFolderQueries[index]?.isPending), ) @@ -395,37 +387,37 @@ export function AgentWorkingDirectoryPanel({ const fileReadQueryOptions = source.type === 'agent' ? consoleQuery.agent.byAgentId.sandbox.files.read.get.queryOptions({ - input: - source.conversationId && selectedWorkingDirectoryFile?.id - ? { - params: { - agent_id: source.agentId, - }, - query: { - conversation_id: source.conversationId, - path: toSandboxApiPath(selectedWorkingDirectoryFile.id), - }, - } - : skipToken, + input: selectedWorkingDirectoryFile?.id + ? { + params: { + agent_id: source.agentId, + }, + query: { + caller_type: source.callerType, + caller_id: source.callerId, + path: toSandboxApiPath(selectedWorkingDirectoryFile.id), + }, + } + : skipToken, context: { silent: true, }, }) : consoleQuery.apps.byAppId.workflowRuns.byWorkflowRunId.agentNodes.byNodeId.sandbox.files.read.get.queryOptions( { - input: - source.appId && workflowNodeRunId && selectedWorkingDirectoryFile?.id - ? { - params: { - app_id: source.appId, - workflow_run_id: workflowNodeRunId, - node_id: source.nodeId, - }, - query: { - path: toSandboxApiPath(selectedWorkingDirectoryFile.id), - }, - } - : skipToken, + input: selectedWorkingDirectoryFile?.id + ? { + params: { + app_id: source.appId, + workflow_run_id: source.workflowRunId, + node_id: source.nodeId, + }, + query: { + node_execution_id: source.nodeExecutionId, + path: toSandboxApiPath(selectedWorkingDirectoryFile.id), + }, + } + : skipToken, context: { silent: true, }, @@ -474,8 +466,9 @@ export function AgentWorkingDirectoryPanel({ 'image-preview', source.type, source.type === 'agent' ? source.agentId : source.appId, - source.type === 'agent' ? source.conversationId : workflowNodeRunId, - source.type === 'workflow-node' ? source.nodeId : undefined, + source.type === 'agent' ? source.callerType : source.workflowRunId, + source.type === 'agent' ? source.callerId : source.nodeId, + source.type === 'workflow-node' ? source.nodeExecutionId : undefined, selectedWorkingDirectoryFilePath, ], queryFn: async () => { @@ -483,44 +476,39 @@ export function AgentWorkingDirectoryPanel({ throw new Error('Missing selected working directory file') if (source.type === 'agent') { - if (!source.conversationId) throw new Error('Missing agent sandbox conversation ID') - return consoleClient.agent.byAgentId.sandbox.files.upload.post({ params: { agent_id: source.agentId, }, body: { - conversation_id: source.conversationId, + caller_type: source.callerType, + caller_id: source.callerId, path: toSandboxApiPath(selectedWorkingDirectoryFilePath), }, }) } - if (!source.appId || !workflowNodeRunId) throw new Error('Missing workflow sandbox source') - return consoleClient.apps.byAppId.workflowRuns.byWorkflowRunId.agentNodes.byNodeId.sandbox.files.upload.post( { params: { app_id: source.appId, - workflow_run_id: workflowNodeRunId, + workflow_run_id: source.workflowRunId, node_id: source.nodeId, }, body: { + node_execution_id: source.nodeExecutionId, path: toSandboxApiPath(selectedWorkingDirectoryFilePath), }, }, ) }, - enabled: - open && !!selectedWorkingDirectoryFile && isImagePreviewFile && hasWorkingDirectorySource, + enabled: open && !!selectedWorkingDirectoryFile && isImagePreviewFile, }) const handleDownloadFile = useCallback( async (action: AgentSkillDetailDownloadAction) => { if (!selectedWorkingDirectoryFile || isFileDownloadPending) return if (source.type === 'agent') { - if (!source.conversationId) return - setDownloadActionLoadingTarget(action) try { const result = await uploadAgentSandboxFile({ @@ -528,7 +516,8 @@ export function AgentWorkingDirectoryPanel({ agent_id: source.agentId, }, body: { - conversation_id: source.conversationId, + caller_type: source.callerType, + caller_id: source.callerId, path: toSandboxApiPath(selectedWorkingDirectoryFile.id), }, }) @@ -540,17 +529,16 @@ export function AgentWorkingDirectoryPanel({ return } - if (!source.appId || !workflowNodeRunId) return - setDownloadActionLoadingTarget(action) try { const result = await uploadWorkflowSandboxFile({ params: { app_id: source.appId, - workflow_run_id: workflowNodeRunId, + workflow_run_id: source.workflowRunId, node_id: source.nodeId, }, body: { + node_execution_id: source.nodeExecutionId, path: toSandboxApiPath(selectedWorkingDirectoryFile.id), }, }) @@ -567,7 +555,6 @@ export function AgentWorkingDirectoryPanel({ tCommon, uploadAgentSandboxFile, uploadWorkflowSandboxFile, - workflowNodeRunId, ], ) diff --git a/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts b/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts index 65c974c7ff3..94304e7a0fa 100644 --- a/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts +++ b/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts @@ -252,6 +252,7 @@ export function useAgentConfigureBuildDraftData({ : activeVersionId, agentSoulConfig: visibleAgentSoulConfig, buildDraftAgentSoulConfig, + id: buildDraftData?.draft.id, changedKeys: buildDraftChangeSummary.changedKeys, changeSummary: buildDraftChangeSummary, changesCount: buildDraftChangeSummary.changesCount,