feat(vm): persist containers across sessions

README.md

@@ -13,8 +13,11 @@ conversations can be resumed with context. One example tool is included:
   Execution happens asynchronously so the assistant can continue responding
   while the command runs.
   The VM is created when a chat session starts and reused for all subsequent
-  tool calls. The environment includes Python and ``pip`` so complex tasks can
-  be scripted using Python directly inside the terminal.
+  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.
 
 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
@@ -77,6 +80,9 @@ 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.
 
+Set ``PERSIST_VMS=0`` to revert to the previous behaviour where containers are
+stopped once no sessions are using them.
+
 To use a fully featured Ubuntu environment, build a custom Docker image and set
 ``VM_IMAGE`` to that image. An example ``docker/Dockerfile.vm`` is provided:
 

run.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 
 from src.chat import ChatSession
+from src.vm import VMRegistry
 
 
 async def _main() -> None:
@@ -24,3 +25,5 @@ if __name__ == "__main__":
         asyncio.run(_main())
     except KeyboardInterrupt:
         pass
+    finally:
+        VMRegistry.shutdown_all()

src/config.py

@@ -10,6 +10,7 @@ MAX_TOOL_CALL_DEPTH: Final[int] = 5
 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"
 
 SYSTEM_PROMPT: Final[str] = (
     "You are Starlette, a professional AI assistant with advanced tool orchestration. "

src/vm.py

@@ -8,7 +8,7 @@ from pathlib import Path
 
 from threading import Lock
 
-from .config import UPLOAD_DIR, VM_IMAGE
+from .config import UPLOAD_DIR, VM_IMAGE, PERSIST_VMS
 
 from .log import get_logger
 
@@ -180,6 +180,18 @@ class VMRegistry:
 
             cls._counts[username] -= 1
             if cls._counts[username] <= 0:
+                cls._counts[username] = 0
+                if not PERSIST_VMS:
+                    vm.stop()
+                    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()
-                del cls._vms[username]
-                del cls._counts[username]
+            cls._vms.clear()
+            cls._counts.clear()