From 6cb8d03bf6301e07033bd20299eea476189f346c Mon Sep 17 00:00:00 2001
From: Harry <xh001x@hotmail.com>
Date: Thu, 15 Jan 2026 00:28:41 +0800
Subject: [PATCH] feat(sandbox): enhance SandboxLayer with app_id handling and
 storage integration

- Introduce _app_id attribute to store application ID from system variables
- Add _get_app_id method to retrieve and validate app_id
- Update on_graph_start to log app_id during sandbox initialization
- Integrate ArchiveSandboxStorage for persisting and restoring sandbox files
- Ensure proper error handling for sandbox file operations
---
 api/core/app/layers/sandbox_layer.py | 41 ++++++++++++++++++++++++----
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/api/core/app/layers/sandbox_layer.py b/api/core/app/layers/sandbox_layer.py
index fa4ec2dc70..aa7d8005f7 100644
--- a/api/core/app/layers/sandbox_layer.py
+++ b/api/core/app/layers/sandbox_layer.py
@@ -1,9 +1,11 @@
 import logging
 
 from core.sandbox.manager import SandboxManager
+from core.sandbox.storage import ArchiveSandboxStorage
 from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
 from core.workflow.graph_engine.layers.base import GraphEngineLayer
 from core.workflow.graph_events.base import GraphEngineEvent
+from extensions.ext_storage import storage
 
 logger = logging.getLogger(__name__)
 
@@ -17,6 +19,7 @@ class SandboxLayer(GraphEngineLayer):
         super().__init__()
         self._tenant_id = tenant_id
         self._workflow_execution_id: str | None = None
+        self._app_id: str | None = None
 
     def _get_workflow_execution_id(self) -> str:
         workflow_execution_id = self.graph_runtime_state.system_variable.workflow_execution_id
@@ -24,6 +27,12 @@ class SandboxLayer(GraphEngineLayer):
             raise RuntimeError("workflow_execution_id is not set in system variables")
         return workflow_execution_id
 
+    def _get_app_id(self) -> str:
+        app_id = self.graph_runtime_state.system_variable.app_id
+        if not app_id:
+            raise RuntimeError("app_id is not set in system variables")
+        return app_id
+
     @property
     def sandbox(self) -> VirtualEnvironment:
         if self._workflow_execution_id is None:
@@ -35,17 +44,18 @@ class SandboxLayer(GraphEngineLayer):
 
     def on_graph_start(self) -> None:
         self._workflow_execution_id = self._get_workflow_execution_id()
-
+        self._app_id = self._get_app_id()
         try:
             from core.sandbox.initializer import AppAssetsInitializer, DifyCliInitializer
             from services.sandbox.sandbox_provider_service import SandboxProviderService
 
-            app_id = self.graph_runtime_state.system_variable.app_id
-            logger.info("Initializing sandbox for tenant_id=%s, app_id=%s", self._tenant_id, app_id)
+            logger.info("Initializing sandbox for tenant_id=%s, app_id=%s", self._tenant_id, self._app_id)
 
-            builder = SandboxProviderService.create_sandbox_builder(self._tenant_id).initializer(DifyCliInitializer())
-            if app_id:
-                builder.initializer(AppAssetsInitializer(self._tenant_id, app_id))
+            builder = (
+                SandboxProviderService.create_sandbox_builder(self._tenant_id)
+                .initializer(DifyCliInitializer())
+                .initializer(AppAssetsInitializer(self._tenant_id, self._app_id))
+            )
             sandbox = builder.build()
 
             SandboxManager.register(self._workflow_execution_id, sandbox)
@@ -55,6 +65,14 @@ class SandboxLayer(GraphEngineLayer):
                 sandbox.metadata.id,
                 sandbox.metadata.arch,
             )
+
+            sandbox_storage = ArchiveSandboxStorage(
+                storage=storage,
+                tenant_id=self._tenant_id,
+                sandbox_id=self._workflow_execution_id,
+            )
+            if sandbox_storage.mount(sandbox):
+                logger.info("Sandbox files restored, workflow_execution_id=%s", self._workflow_execution_id)
         except Exception as e:
             logger.exception("Failed to initialize sandbox")
             raise SandboxInitializationError(f"Failed to initialize sandbox: {e}") from e
@@ -79,6 +97,17 @@ class SandboxLayer(GraphEngineLayer):
             sandbox_id,
         )
 
+        try:
+            sandbox_storage = ArchiveSandboxStorage(
+                storage=storage,
+                tenant_id=self._tenant_id,
+                sandbox_id=self._workflow_execution_id,
+            )
+            sandbox_storage.unmount(sandbox)
+            logger.info("Sandbox files persisted, workflow_execution_id=%s", self._workflow_execution_id)
+        except Exception:
+            logger.exception("Failed to persist sandbox files, workflow_execution_id=%s", self._workflow_execution_id)
+
         try:
             sandbox.release_environment()
             logger.info("Sandbox released, sandbox_id=%s", sandbox_id)