Merge pull request #46 from EnvisionMindCa/codex/make-vm-persistent-with-saved-state

README.md

@@ -14,10 +14,10 @@ conversations can be resumed with context. One example tool is included:
   while the command runs.
   The VM is created when a chat session starts and reused for all subsequent
   tool calls. When ``PERSIST_VMS`` is enabled (default), each user keeps the
-  same container across multiple chat sessions so any installed packages and
-  filesystem changes remain available. The environment includes Python and
-  ``pip`` so complex tasks can be scripted using Python directly inside the
-  terminal.
+  same container across multiple chat sessions and across application restarts,
+  so any installed packages and filesystem changes remain available. The
+  environment includes Python and ``pip`` so complex tasks can be scripted using
+  Python directly inside the terminal.
 
 Sessions share state through an in-memory registry so that only one generation
 can run at a time. Messages sent while a response is being produced are
@@ -80,6 +80,11 @@ back to ``python:3.11-slim``. This base image includes Python and ``pip`` so
 packages can be installed immediately. The container has network access enabled
 which allows fetching additional dependencies as needed.
 
+When ``PERSIST_VMS`` is ``1`` (default), containers are kept around and reused
+across application restarts. Each user is assigned a stable container name, so
+packages installed or files created inside the VM remain available the next
+time the application starts. Set ``VM_STATE_DIR`` to specify the host directory
+used for per-user persistent storage mounted inside the VM at ``/state``.
 Set ``PERSIST_VMS=0`` to revert to the previous behaviour where containers are
 stopped once no sessions are using them.
 

src/config.py

@@ -11,6 +11,9 @@ NUM_CTX: Final[int] = int(os.getenv("OLLAMA_NUM_CTX", "32000"))
 UPLOAD_DIR: Final[str] = os.getenv("UPLOAD_DIR", str(Path.cwd() / "uploads"))
 VM_IMAGE: Final[str] = os.getenv("VM_IMAGE", "python:3.11")
 PERSIST_VMS: Final[bool] = os.getenv("PERSIST_VMS", "1") == "1"
+VM_STATE_DIR: Final[str] = os.getenv(
+    "VM_STATE_DIR", str(Path.cwd() / "vm_state")
+)
 
 SYSTEM_PROMPT: Final[str] = (
     "You are Starlette, a professional AI assistant with advanced tool orchestration. "

src/vm.py

@@ -8,13 +8,19 @@ from pathlib import Path
 
 from threading import Lock
 
-from .config import UPLOAD_DIR, VM_IMAGE, PERSIST_VMS
+from .config import UPLOAD_DIR, VM_IMAGE, PERSIST_VMS, VM_STATE_DIR
 
 from .log import get_logger
 
 _LOG = get_logger(__name__)
 
 
+def _sanitize(name: str) -> str:
+    """Return a Docker-safe name fragment."""
+    allowed = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_.")
+    return "".join(c if c in allowed else "_" for c in name)
+
+
 class LinuxVM:
     """Manage a lightweight Docker-based VM.
 
@@ -23,13 +29,18 @@ class LinuxVM:
     """
 
     def __init__(
-        self, image: str = VM_IMAGE, host_dir: str = UPLOAD_DIR
+        self,
+        username: str,
+        image: str = VM_IMAGE,
+        host_dir: str = UPLOAD_DIR,
     ) -> None:
         self._image = image
-        self._name = f"chat-vm-{uuid.uuid4().hex[:8]}"
+        self._name = f"chat-vm-{_sanitize(username)}"
         self._running = False
         self._host_dir = Path(host_dir)
         self._host_dir.mkdir(parents=True, exist_ok=True)
+        self._state_dir = Path(VM_STATE_DIR) / _sanitize(username)
+        self._state_dir.mkdir(parents=True, exist_ok=True)
 
     def start(self) -> None:
         """Start the VM if it is not already running."""
@@ -37,6 +48,25 @@ class LinuxVM:
             return
 
         try:
+            inspect = subprocess.run(
+                ["docker", "inspect", "-f", "{{.State.Running}}", self._name],
+                capture_output=True,
+                text=True,
+            )
+            if inspect.returncode == 0:
+                if inspect.stdout.strip() == "true":
+                    self._running = True
+                    return
+                subprocess.run(
+                    ["docker", "start", self._name],
+                    check=True,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                )
+                self._running = True
+                return
+
             subprocess.run(
                 ["docker", "pull", self._image],
                 check=False,
@@ -53,6 +83,8 @@ class LinuxVM:
                     self._name,
                     "-v",
                     f"{self._host_dir}:/data",
+                    "-v",
+                    f"{self._state_dir}:/state",
                     self._image,
                     "sleep",
                     "infinity",
@@ -130,13 +162,22 @@ class LinuxVM:
         if not self._running:
             return
 
-        subprocess.run(
-            ["docker", "rm", "-f", self._name],
-            check=False,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-        )
+        if PERSIST_VMS:
+            subprocess.run(
+                ["docker", "stop", self._name],
+                check=False,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+        else:
+            subprocess.run(
+                ["docker", "rm", "-f", self._name],
+                check=False,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
         self._running = False
 
     def __enter__(self) -> "LinuxVM":
@@ -161,7 +202,10 @@ class VMRegistry:
         with cls._lock:
             vm = cls._vms.get(username)
             if vm is None:
-                vm = LinuxVM(host_dir=str(Path(UPLOAD_DIR) / username))
+                vm = LinuxVM(
+                    username,
+                    host_dir=str(Path(UPLOAD_DIR) / username),
+                )
                 cls._vms[username] = vm
                 cls._counts[username] = 0
             cls._counts[username] += 1
@@ -183,15 +227,16 @@ class VMRegistry:
                 cls._counts[username] = 0
                 if not PERSIST_VMS:
                     vm.stop()
-                    del cls._vms[username]
-                    del cls._counts[username]
+                del cls._vms[username]
+                del cls._counts[username]
 
     @classmethod
     def shutdown_all(cls) -> None:
         """Stop and remove all managed VMs."""
 
         with cls._lock:
-            for vm in cls._vms.values():
-                vm.stop()
+            if not PERSIST_VMS:
+                for vm in cls._vms.values():
+                    vm.stop()
             cls._vms.clear()
             cls._counts.clear()