rewrite snapshot script: use run_commands, dotenv, idempotent delete

surfsense_backend/scripts/create_sandbox_snapshot.py

@@ -1,16 +1,32 @@
-"""Create the Daytona snapshot used by SurfSense sandboxes.
+"""Create the Daytona snapshot used by SurfSense code-execution sandboxes.
 
-Usage:
+Run from the backend directory:
+    cd surfsense_backend
     uv run python scripts/create_sandbox_snapshot.py
 
-Requires DAYTONA_API_KEY (and optionally DAYTONA_API_URL / DAYTONA_TARGET)
-to be set in the environment or in a .env file.
+Prerequisites:
+    - DAYTONA_API_KEY set in surfsense_backend/.env (or exported in shell)
+    - DAYTONA_API_URL=https://app.daytona.io/api
+    - DAYTONA_TARGET=us  (or eu)
+
+After this script succeeds, add to surfsense_backend/.env:
+    DAYTONA_SNAPSHOT_ID=surfsense-sandbox
 """
 
 import os
 import sys
+import time
+from pathlib import Path
 
-from daytona import CreateSnapshotParams, Daytona, DaytonaConfig, Image
+from dotenv import load_dotenv
+
+_here = Path(__file__).parent
+for candidate in [_here / "../surfsense_backend/.env", _here / ".env", _here / "../.env"]:
+    if candidate.exists():
+        load_dotenv(candidate)
+        break
+
+from daytona import CreateSnapshotParams, Daytona, Image  # noqa: E402
 
 SNAPSHOT_NAME = "surfsense-sandbox"
 
@@ -23,38 +39,58 @@ PACKAGES = [
 ]
 
 
-def main() -> None:
-    config = DaytonaConfig(
-        api_key=os.environ.get("DAYTONA_API_KEY", ""),
-        api_url=os.environ.get("DAYTONA_API_URL", "https://app.daytona.io/api"),
-        target=os.environ.get("DAYTONA_TARGET", "us"),
-    )
-    daytona = Daytona(config)
-
-    image = (
+def build_image() -> Image:
+    """Build the sandbox image with data-science packages and a /documents symlink."""
+    return (
         Image.debian_slim("3.12")
         .pip_install(*PACKAGES)
-        # The agent's virtual filesystem serves documents at /documents/.
-        # This symlink lets code inside the sandbox use the same path.
-        .run("mkdir -p /home/daytona/documents && ln -sf /home/daytona/documents /documents")
+        # Symlink /documents → /home/daytona/documents so the LLM can use
+        # the same /documents/ path it sees in the virtual filesystem.
+        .run_commands(
+            "mkdir -p /home/daytona/documents",
+            "ln -sfn /home/daytona/documents /documents",
+        )
     )
 
+
+def main() -> None:
+    api_key = os.environ.get("DAYTONA_API_KEY")
+    if not api_key:
+        print("ERROR: DAYTONA_API_KEY is not set.", file=sys.stderr)
+        print("Add it to surfsense_backend/.env or export it in your shell.", file=sys.stderr)
+        sys.exit(1)
+
+    daytona = Daytona()
+
     try:
         existing = daytona.snapshot.get(SNAPSHOT_NAME)
-        print(f"Deleting existing snapshot '{SNAPSHOT_NAME}'...")
+        print(f"Deleting existing snapshot '{SNAPSHOT_NAME}' …")
         daytona.snapshot.delete(existing)
-        print("Deleted.")
+        print(f"Deleted '{SNAPSHOT_NAME}'. Waiting for removal to propagate …")
+        for attempt in range(30):
+            time.sleep(2)
+            try:
+                daytona.snapshot.get(SNAPSHOT_NAME)
+            except Exception:
+                print(f"Confirmed '{SNAPSHOT_NAME}' is gone.\n")
+                break
+        else:
+            print(f"WARNING: '{SNAPSHOT_NAME}' may still exist after 60s. Proceeding anyway.\n")
     except Exception:
         pass
 
-    print(f"Creating snapshot '{SNAPSHOT_NAME}' with packages: {', '.join(PACKAGES)}")
-    snapshot = daytona.snapshot.create(
-        CreateSnapshotParams(name=SNAPSHOT_NAME, image=image),
-        on_logs=lambda chunk: print(chunk, end=""),
+    print(f"Building snapshot '{SNAPSHOT_NAME}' …")
+    print(f"Packages: {', '.join(PACKAGES)}\n")
+
+    daytona.snapshot.create(
+        CreateSnapshotParams(name=SNAPSHOT_NAME, image=build_image()),
+        on_logs=lambda chunk: print(chunk, end="", flush=True),
     )
-    print(f"\nSnapshot created: {snapshot.name}")
-    print(f"Set DAYTONA_SNAPSHOT_ID={snapshot.name} in your .env")
+
+    print(f"\n\nSnapshot '{SNAPSHOT_NAME}' is ready.")
+    print("\nAdd this to surfsense_backend/.env:")
+    print(f"    DAYTONA_SNAPSHOT_ID={SNAPSHOT_NAME}")
 
 
 if __name__ == "__main__":
-    sys.exit(main() or 0)
+    main()