From 226e6c826309c2ab0c92055e6a3c5c3a3dfd6dce Mon Sep 17 00:00:00 2001
From: fuyua9 <fuyua9@foxmail.com>
Date: Sat, 2 May 2026 06:51:34 +0800
Subject: [PATCH] fix(daytona): make sandbox snapshot configurable (#1220)

---
 CUSTOMIZATION.md                  |  2 +-
 agent/integrations/daytona.py     | 13 ++++--
 tests/test_daytona_integration.py | 70 +++++++++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_daytona_integration.py

diff --git a/CUSTOMIZATION.md b/CUSTOMIZATION.md
index 4e45a2c9..91ef4af4 100644
--- a/CUSTOMIZATION.md
+++ b/CUSTOMIZATION.md
@@ -49,7 +49,7 @@ Set the `SANDBOX_TYPE` environment variable to switch providers. Each provider h
 | `SANDBOX_TYPE` | Integration file | Required env vars |
 |---|---|---|
 | `langsmith` (default) | `agent/integrations/langsmith.py` | `LANGSMITH_API_KEY_PROD`, `SANDBOX_TYPE="langsmith"` |
-| `daytona` | `agent/integrations/daytona.py` | `DAYTONA_API_KEY`, `SANDBOX_TYPE="daytona"` |
+| `daytona` | `agent/integrations/daytona.py` | `DAYTONA_API_KEY`, `SANDBOX_TYPE="daytona"`, optional `DAYTONA_SANDBOX_SNAPSHOT` |
 | `runloop` | `agent/integrations/runloop.py` | `RUNLOOP_API_KEY`, `SANDBOX_TYPE="runloop"` |
 | `modal` | `agent/integrations/modal.py` | Modal credentials, `SANDBOX_TYPE="modal"` |
 | `local` | `agent/integrations/local.py` | None (no isolation — development only), `SANDBOX_TYPE="local"` |
diff --git a/agent/integrations/daytona.py b/agent/integrations/daytona.py
index 1482ee86..8a27db89 100644
--- a/agent/integrations/daytona.py
+++ b/agent/integrations/daytona.py
@@ -3,8 +3,15 @@ import os
 from daytona import CreateSandboxFromSnapshotParams, Daytona, DaytonaConfig
 from langchain_daytona import DaytonaSandbox
 
-# TODO: Update this to include your specific sandbox configuration
-DAYTONA_SANDBOX_PARAMS = CreateSandboxFromSnapshotParams(snapshot="daytonaio/sandbox:0.6.0")
+DEFAULT_DAYTONA_SANDBOX_SNAPSHOT = "daytonaio/sandbox:0.6.0"
+DAYTONA_SANDBOX_SNAPSHOT_ENV = "DAYTONA_SANDBOX_SNAPSHOT"
+
+
+def _get_daytona_sandbox_params() -> CreateSandboxFromSnapshotParams:
+    snapshot = os.getenv(DAYTONA_SANDBOX_SNAPSHOT_ENV, DEFAULT_DAYTONA_SANDBOX_SNAPSHOT).strip()
+    if not snapshot:
+        raise ValueError(f"{DAYTONA_SANDBOX_SNAPSHOT_ENV} must not be empty")
+    return CreateSandboxFromSnapshotParams(snapshot=snapshot)
 
 
 def create_daytona_sandbox(sandbox_id: str | None = None):
@@ -17,6 +24,6 @@ def create_daytona_sandbox(sandbox_id: str | None = None):
     if sandbox_id:
         sandbox = daytona.get(sandbox_id)
     else:
-        sandbox = daytona.create(params=DAYTONA_SANDBOX_PARAMS)
+        sandbox = daytona.create(params=_get_daytona_sandbox_params())
 
     return DaytonaSandbox(sandbox=sandbox)
diff --git a/tests/test_daytona_integration.py b/tests/test_daytona_integration.py
new file mode 100644
index 00000000..05d14aa2
--- /dev/null
+++ b/tests/test_daytona_integration.py
@@ -0,0 +1,70 @@
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+class _FakeCreateSandboxFromSnapshotParams:
+    def __init__(self, *, snapshot: str):
+        self.snapshot = snapshot
+
+
+class _FakeDaytonaConfig:
+    def __init__(self, *, api_key: str):
+        self.api_key = api_key
+
+
+class _FakeDaytonaSandbox:
+    def __init__(self, *, sandbox):
+        self.sandbox = sandbox
+
+
+def _load_daytona_module(monkeypatch):
+    fake_daytona = types.ModuleType("daytona")
+    fake_daytona.CreateSandboxFromSnapshotParams = _FakeCreateSandboxFromSnapshotParams
+    fake_daytona.DaytonaConfig = _FakeDaytonaConfig
+    fake_daytona.Daytona = object
+
+    fake_langchain_daytona = types.ModuleType("langchain_daytona")
+    fake_langchain_daytona.DaytonaSandbox = _FakeDaytonaSandbox
+
+    monkeypatch.setitem(sys.modules, "daytona", fake_daytona)
+    monkeypatch.setitem(sys.modules, "langchain_daytona", fake_langchain_daytona)
+    module_path = ROOT / "agent" / "integrations" / "daytona.py"
+    spec = importlib.util.spec_from_file_location("daytona_under_test", module_path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_daytona_params_default_to_existing_snapshot(monkeypatch):
+    monkeypatch.delenv("DAYTONA_SANDBOX_SNAPSHOT", raising=False)
+    module = _load_daytona_module(monkeypatch)
+
+    params = module._get_daytona_sandbox_params()
+
+    assert params.snapshot == "daytonaio/sandbox:0.6.0"
+
+
+def test_daytona_params_use_env_snapshot(monkeypatch):
+    monkeypatch.setenv("DAYTONA_SANDBOX_SNAPSHOT", "custom/snapshot:1.0")
+    module = _load_daytona_module(monkeypatch)
+
+    params = module._get_daytona_sandbox_params()
+
+    assert params.snapshot == "custom/snapshot:1.0"
+
+
+def test_daytona_params_reject_empty_snapshot(monkeypatch):
+    monkeypatch.setenv("DAYTONA_SANDBOX_SNAPSHOT", "  ")
+    module = _load_daytona_module(monkeypatch)
+
+    try:
+        module._get_daytona_sandbox_params()
+    except ValueError as exc:
+        assert "DAYTONA_SANDBOX_SNAPSHOT must not be empty" in str(exc)
+    else:
+        raise AssertionError("expected empty Daytona snapshot to fail")