Add per-user VM registry and update upload paths

src/chat.py

@@ -27,7 +27,7 @@ from .db import (
 from .log import get_logger
 from .schema import Msg
 from .tools import execute_terminal, set_vm
-from .vm import LinuxVM
+from .vm import VMRegistry
 
 _LOG = get_logger(__name__)
 
@@ -48,18 +48,19 @@ class ChatSession:
         self._conversation, _ = Conversation.get_or_create(
             user=self._user, session_name=session
         )
-        self._vm = LinuxVM()
+        self._vm = None
         self._messages: List[Msg] = self._load_history()
         self._ensure_system_prompt()
 
     async def __aenter__(self) -> "ChatSession":
-        self._vm.start()
+        self._vm = VMRegistry.acquire(self._user.username)
         set_vm(self._vm)
         return self
 
     async def __aexit__(self, exc_type, exc, tb) -> None:
         set_vm(None)
-        self._vm.stop()
+        if self._vm:
+            VMRegistry.release(self._user.username)
         if not _db.is_closed():
             _db.close()
 
@@ -79,7 +80,7 @@ class ChatSession:
         target = dest / src.name
         shutil.copy(src, target)
         add_document(self._user.username, str(target), src.name)
-        return f"/data/{self._user.username}/{src.name}"
+        return f"/data/{src.name}"
 
     def _ensure_system_prompt(self) -> None:
         if any(m.get("role") == "system" for m in self._messages):

src/vm.py

@@ -4,6 +4,8 @@ import subprocess
 import uuid
 from pathlib import Path
 
+from threading import Lock
+
 from .config import UPLOAD_DIR
 
 from .log import get_logger
@@ -108,3 +110,41 @@ class LinuxVM:
 
     def __exit__(self, exc_type, exc, tb) -> None:
         self.stop()
+
+
+class VMRegistry:
+    """Manage Linux VM instances on a per-user basis."""
+
+    _vms: dict[str, LinuxVM] = {}
+    _counts: dict[str, int] = {}
+    _lock = Lock()
+
+    @classmethod
+    def acquire(cls, username: str) -> LinuxVM:
+        """Return a running VM for ``username``, creating it if needed."""
+
+        with cls._lock:
+            vm = cls._vms.get(username)
+            if vm is None:
+                vm = LinuxVM(host_dir=str(Path(UPLOAD_DIR) / username))
+                cls._vms[username] = vm
+                cls._counts[username] = 0
+            cls._counts[username] += 1
+
+        vm.start()
+        return vm
+
+    @classmethod
+    def release(cls, username: str) -> None:
+        """Release one reference to ``username``'s VM and stop it if unused."""
+
+        with cls._lock:
+            vm = cls._vms.get(username)
+            if vm is None:
+                return
+
+            cls._counts[username] -= 1
+            if cls._counts[username] <= 0:
+                vm.stop()
+                del cls._vms[username]
+                del cls._counts[username]